diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java new file mode 100644 index 000000000000..39b562fe0f8f --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java @@ -0,0 +1,138 @@ +/* + * 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.io.streaming; + +import java.util.Arrays; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.util.SerializableConfiguration; + +/** One split of a Beam unbounded source for one micro-batch, from epoch start to epoch end. */ +public class BeamInputPartition implements InputPartition { + + private static final long serialVersionUID = 1L; + + private final UnboundedSource split; + private final Coder> coder; + private final Broadcast options; + private final Broadcast hadoopConf; + private final String checkpointLocation; + private final int splitId; + private final long startEpoch; + private final long endEpoch; + private final long maxRecords; + private final long maxBatchDurationMillis; + private final long readerIdleTimeoutMillis; + private final String[] preferredLocations; + + BeamInputPartition( + UnboundedSource split, + Coder> coder, + Broadcast options, + Broadcast hadoopConf, + String checkpointLocation, + int splitId, + long startEpoch, + long endEpoch, + long maxRecords, + long maxBatchDurationMillis, + long readerIdleTimeoutMillis, + String[] preferredLocations) { + this.split = split; + this.coder = coder; + this.options = options; + this.hadoopConf = hadoopConf; + this.checkpointLocation = checkpointLocation; + this.splitId = splitId; + this.startEpoch = startEpoch; + this.endEpoch = endEpoch; + this.maxRecords = maxRecords; + this.maxBatchDurationMillis = maxBatchDurationMillis; + this.readerIdleTimeoutMillis = readerIdleTimeoutMillis; + this.preferredLocations = preferredLocations.clone(); + } + + UnboundedSource split() { + return split; + } + + Coder> coder() { + return coder; + } + + Broadcast options() { + return options; + } + + Broadcast hadoopConf() { + return hadoopConf; + } + + String checkpointLocation() { + return checkpointLocation; + } + + int splitId() { + return splitId; + } + + long startEpoch() { + return startEpoch; + } + + long endEpoch() { + return endEpoch; + } + + /** Records this split may emit in this micro-batch, below 1 means unlimited. */ + long maxRecords() { + return maxRecords; + } + + long maxBatchDurationMillis() { + return maxBatchDurationMillis; + } + + long readerIdleTimeoutMillis() { + return readerIdleTimeoutMillis; + } + + @Override + public String[] preferredLocations() { + return preferredLocations.clone(); + } + + @Override + public String toString() { + return "BeamInputPartition{checkpointLocation=" + + checkpointLocation + + ", split=" + + splitId + + ", epochs=" + + startEpoch + + ".." + + endEpoch + + ", locations=" + + Arrays.toString(preferredLocations) + + "}"; + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java new file mode 100644 index 000000000000..0ec8fb061f4e --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java @@ -0,0 +1,277 @@ +/* + * 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.io.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.spark.SparkEnv; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.connector.read.streaming.Offset; +import org.apache.spark.storage.BlockManager; +import org.apache.spark.storage.BlockManagerId; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import scala.collection.Iterator; + +/** + * Driver side {@link MicroBatchStream} over a Beam {@link UnboundedSource}. + * + *

Offsets are opaque epochs, {@link #latestOffset()} advances by one every trigger. The source + * is split once and the splits are pinned under the checkpoint location, every batch of every run + * plans the same partitions. {@link #commit} purges marks below the committed epoch on a background + * thread, Spark never asks for those again. + * + *

Each split prefers the executor rendezvous hashing assigns it, so an executor joining or + * leaving moves only the splits of that executor. Locality is a hint, the reader cache restores a + * split from its durable mark wherever it lands. + */ +public class BeamMicroBatchStream implements MicroBatchStream { + + private static final Logger LOG = LoggerFactory.getLogger(BeamMicroBatchStream.class); + + private final BeamSourceSpec spec; + private final String checkpointLocation; + private final BeamSourceCheckpoint checkpoint; + + private final ExecutorService purger = + Executors.newSingleThreadExecutor( + runnable -> { + Thread thread = new Thread(runnable, "beam-source-mark-purge"); + thread.setDaemon(true); + return thread; + }); + private final AtomicBoolean purgeInFlight = new AtomicBoolean(); + private final AtomicLong purgeRequested = new AtomicLong(); + + private long epoch; + private @Nullable List> splits; + + BeamMicroBatchStream(BeamSourceSpec spec, String checkpointLocation) { + this.spec = spec; + this.checkpointLocation = checkpointLocation; + this.checkpoint = + new BeamSourceCheckpoint(checkpointLocation, spec.hadoopConf().value().value()); + } + + @Override + public Offset initialOffset() { + return BeamOffset.ZERO; + } + + @Override + public synchronized Offset latestOffset() { + return new BeamOffset(++epoch); + } + + @Override + public Offset deserializeOffset(String json) { + BeamOffset offset = BeamOffset.fromJson(json); + fastForwardEpoch(offset.epoch()); + return offset; + } + + @Override + public InputPartition[] planInputPartitions(Offset start, Offset end) { + long startEpoch = ((BeamOffset) start).epoch(); + long endEpoch = ((BeamOffset) end).epoch(); + fastForwardEpoch(endEpoch); + List> pinned = splits(); + long[] quotas = splitQuotas(spec.maxRecordsPerBatch(), pinned.size()); + List executors = sortedExecutors(); + InputPartition[] partitions = new InputPartition[pinned.size()]; + for (int i = 0; i < pinned.size(); i++) { + String[] locations = + executors.isEmpty() ? new String[0] : new String[] {assign(i, executors)}; + partitions[i] = + new BeamInputPartition<>( + pinned.get(i), + spec.coder(), + spec.options(), + spec.hadoopConf(), + checkpointLocation, + i, + startEpoch, + endEpoch, + quotas[i], + spec.maxBatchDurationMillis(), + spec.readerIdleTimeoutMillis(), + locations); + } + return partitions; + } + + @Override + public PartitionReaderFactory createReaderFactory() { + return new BeamPartitionReaderFactory(); + } + + /** Purges marks below {@code end} off the stream thread, one purge runs at a time. */ + @Override + public void commit(Offset end) { + long endEpoch = ((BeamOffset) end).epoch(); + int numSplits = splits().size(); + purgeRequested.accumulateAndGet(endEpoch, Math::max); + if (purgeInFlight.compareAndSet(false, true)) { + purger.execute(() -> purgeRequested(numSplits)); + } + } + + private void purgeRequested(int numSplits) { + long epoch; + do { + epoch = purgeRequested.get(); + try { + for (int i = 0; i < numSplits; i++) { + checkpoint.purgeMarksBelow(i, epoch); + } + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to purge marks below epoch {} at {}.", epoch, checkpointLocation, e); + } + purgeInFlight.set(false); + } while (purgeRequested.get() > epoch && purgeInFlight.compareAndSet(false, true)); + } + + @Override + public void stop() { + LOG.info( + "Stopping Beam micro-batch stream {} at {}.", spec.transformName(), checkpointLocation); + purger.shutdown(); + } + + /** Keeps {@link #latestOffset()} ahead of every epoch Spark logged before a restart. */ + private synchronized void fastForwardEpoch(long seen) { + if (seen > epoch) { + LOG.info("Fast forwarding epoch of {} from {} to {}.", spec.transformName(), epoch, seen); + epoch = seen; + } + } + + private synchronized List> splits() { + if (splits != null) { + return splits; + } + List> pinned; + try { + pinned = checkpoint.readSplits(); + } catch (IOException e) { + throw new IllegalStateException("Failed to read pinned splits at " + checkpointLocation, e); + } + if (pinned == null) { + pinned = new ArrayList<>(splitSource()); + try { + checkpoint.writeSplits(pinned); + } catch (IOException e) { + throw new IllegalStateException("Failed to pin splits at " + checkpointLocation, e); + } + } else { + LOG.info("Restored {} pinned split(s) from {}.", pinned.size(), checkpointLocation); + } + List> typed = new ArrayList<>(pinned.size()); + for (UnboundedSource split : pinned) { + @SuppressWarnings("unchecked") // splits of this source share its element type + UnboundedSource cast = (UnboundedSource) split; + typed.add(cast); + } + splits = typed; + return typed; + } + + private List> splitSource() { + UnboundedSource source = spec.source(); + PipelineOptions options = spec.options().value().get(); + List> result; + try { + result = source.split(spec.desiredNumSplits(), options); + } catch (Exception e) { + throw new IllegalStateException( + "Failed to split UnboundedSource " + source.getClass().getCanonicalName(), e); + } + if (result.isEmpty()) { + result = Collections.singletonList(source); + } + LOG.info( + "Split {} into {} partition(s), desired {}.", + spec.transformName(), + result.size(), + spec.desiredNumSplits()); + return result; + } + + /** + * Divides the batch quota over the splits, the remainder goes to the first splits. Below 1 means + * unlimited for every split. A quota below the split count gives every split one record. + */ + static long[] splitQuotas(long maxRecordsPerBatch, int numSplits) { + long[] quotas = new long[numSplits]; + if (maxRecordsPerBatch < 1) { + Arrays.fill(quotas, maxRecordsPerBatch); + return quotas; + } + long base = maxRecordsPerBatch / numSplits; + long remainder = maxRecordsPerBatch % numSplits; + for (int i = 0; i < numSplits; i++) { + quotas[i] = Math.max(1L, base + (i < remainder ? 1 : 0)); + } + return quotas; + } + + /** Rendezvous hashing of a split over the executors, stable under membership changes. */ + static String assign(int splitId, List executors) { + String best = executors.get(0); + int bestHash = Integer.MIN_VALUE; + for (String executor : executors) { + int hash = Hashing.murmur3_32_fixed().hashUnencodedChars(executor + '#' + splitId).asInt(); + if (hash > bestHash) { + bestHash = hash; + best = executor; + } + } + return best; + } + + /** Sorted {@code executor__} locations, empty in local mode or on any failure. */ + private static List sortedExecutors() { + try { + BlockManager bm = SparkEnv.get().blockManager(); + Iterator peers = bm.master().getPeers(bm.blockManagerId()).iterator(); + List executors = new ArrayList<>(); + while (peers.hasNext()) { + BlockManagerId peer = peers.next(); + executors.add("executor_" + peer.host() + "_" + peer.executorId()); + } + Collections.sort(executors); + return executors; + } catch (RuntimeException e) { + LOG.debug("No executor list available for preferred locations.", e); + return Collections.emptyList(); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java new file mode 100644 index 000000000000..ba5e9825e27a --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java @@ -0,0 +1,59 @@ +/* + * 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.io.streaming; + +import org.apache.spark.sql.connector.read.streaming.Offset; + +/** + * Opaque epoch counter used as the Spark {@link Offset} of a Beam unbounded source. + * + *

The read position lives in Beam checkpoint marks on the executors, see {@link + * BeamSourceCheckpoint}. Equality is the base class comparison of {@link #json()}. + */ +public class BeamOffset extends Offset { + + public static final BeamOffset ZERO = new BeamOffset(0L); + + private final long epoch; + + public BeamOffset(long epoch) { + this.epoch = epoch; + } + + public long epoch() { + return epoch; + } + + @Override + public String json() { + return Long.toString(epoch); + } + + public static BeamOffset fromJson(String json) { + try { + return new BeamOffset(Long.parseLong(json.trim())); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Not a valid BeamOffset: " + json, e); + } + } + + @Override + public String toString() { + return json(); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java new file mode 100644 index 000000000000..d500aec44644 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -0,0 +1,209 @@ +/* + * 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.io.streaming; + +import java.io.IOException; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache.CachedReader; +import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.util.BackOff; +import org.apache.beam.sdk.util.BackOffUtils; +import org.apache.beam.sdk.util.FluentBackoff; +import org.apache.beam.sdk.util.Sleeper; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.TaskContext; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Reads one split of a Beam {@link UnboundedSource} for one micro-batch. + * + *

The batch ends at the record quota or at the deadline. The reader then writes its checkpoint + * mark durably at the end epoch and stays in {@link BeamReaderCache} for the next batch. A failed + * mark write fails the task. An attempt Spark killed or failed writes nothing and its reader is + * dropped, the retry restores from the durable mark at the start epoch. + * + * @param the element type of the split + */ +public class BeamPartitionReader implements PartitionReader { + + private static final Logger LOG = LoggerFactory.getLogger(BeamPartitionReader.class); + + private static final Duration INITIAL_BACKOFF = Duration.millis(10); + + private final String key; + private final UnboundedSource split; + private final Coder> coder; + private final BeamSourceCheckpoint checkpoint; + private final CachedReader cached; + private final int splitId; + private final long endEpoch; + private final long maxRecords; + private final long maxBatchDurationMillis; + + private long recordsRead; + private long deadlineMillis = -1L; + private boolean batchEnded; + private @Nullable InternalRow current; + + BeamPartitionReader(BeamInputPartition partition) throws IOException { + this.split = partition.split(); + this.coder = partition.coder(); + this.splitId = partition.splitId(); + this.endEpoch = partition.endEpoch(); + this.maxRecords = partition.maxRecords(); + this.maxBatchDurationMillis = partition.maxBatchDurationMillis(); + PipelineOptions options = partition.options().value().get(); + Configuration conf = partition.hadoopConf().value().value(); + this.checkpoint = new BeamSourceCheckpoint(partition.checkpointLocation(), conf); + this.key = BeamReaderCache.key(partition.checkpointLocation(), splitId); + long startEpoch = partition.startEpoch(); + this.cached = + BeamReaderCache.acquire( + key, + startEpoch, + split, + options, + partition.readerIdleTimeoutMillis(), + () -> checkpoint.readMark(splitId, startEpoch)); + } + + @Override + public boolean next() throws IOException { + if (deadlineMillis < 0) { + deadlineMillis = System.currentTimeMillis() + maxBatchDurationMillis; + } + BackOff backOff = null; + while (true) { + if (maxRecords > 0 && recordsRead >= maxRecords) { + return endOfBatch(false); + } + long remaining = deadlineMillis - System.currentTimeMillis(); + if (remaining <= 0) { + return endOfBatch(false); + } + if (cached.startOrAdvance()) { + recordsRead++; + current = toRow(); + return true; + } + if (backOff == null) { + backOff = backOff(remaining); + } + try { + if (!BackOffUtils.next(Sleeper.DEFAULT, backOff)) { + return endOfBatch(false); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return endOfBatch(true); + } + } + } + + @Override + public InternalRow get() { + InternalRow row = current; + if (row == null) { + throw new IllegalStateException("No current row, next() did not return true."); + } + return row; + } + + @Override + public void close() throws IOException { + endBatch(attemptDiscarded()); + current = null; + } + + private boolean endOfBatch(boolean discarded) throws IOException { + endBatch(discarded); + current = null; + return false; + } + + /** + * Ends the batch once. A discarded attempt drops the reader and writes nothing. A reader that was + * never started has not moved, its start mark is written forward, an empty file standing for a + * fresh start. + */ + private void endBatch(boolean discarded) throws IOException { + if (batchEnded) { + return; + } + batchEnded = true; + if (discarded) { + LOG.info("Attempt for Beam reader {} was discarded, dropping the reader.", key); + BeamReaderCache.invalidate(key); + return; + } + if (!cached.started()) { + byte[] startMark = cached.positionMark(); + byte[] codedMark = startMark == null ? new byte[0] : startMark; + checkpoint.writeMark(splitId, endEpoch, codedMark); + cached.endBatch(endEpoch, null, codedMark); + return; + } + CheckpointMark mark = cached.reader().getCheckpointMark(); + byte[] codedMark = encodeMark(split, mark); + checkpoint.writeMark(splitId, endEpoch, codedMark); + cached.endBatch(endEpoch, mark, codedMark); + LOG.debug("Beam reader {} read {} record(s) up to epoch {}.", key, recordsRead, endEpoch); + } + + private static boolean attemptDiscarded() { + TaskContext context = TaskContext.get(); + return context != null && (context.isInterrupted() || context.isFailed()); + } + + private static byte[] encodeMark( + UnboundedSource source, CheckpointMark mark) { + @SuppressWarnings("unchecked") // getCheckpointMark returns the source's own mark type + MarkT typed = (MarkT) mark; + return CoderHelpers.toByteArray(typed, source.getCheckpointMarkCoder()); + } + + private static BackOff backOff(long remainingMillis) { + Duration remaining = Duration.millis(remainingMillis); + return FluentBackoff.DEFAULT + .withInitialBackoff(INITIAL_BACKOFF) + .withMaxBackoff(remaining) + .withMaxCumulativeBackoff(remaining) + .backoff(); + } + + private InternalRow toRow() { + Instant timestamp = cached.reader().getCurrentTimestamp(); + WindowedValue value = + WindowedValues.timestampedValueInGlobalWindow(cached.reader().getCurrent(), timestamp); + byte[] payload = CoderHelpers.toByteArray(value, coder); + // Spark stores TimestampType as microseconds. + return new GenericInternalRow(new Object[] {payload, timestamp.getMillis() * 1000L}); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java new file mode 100644 index 000000000000..a360cf895b7b --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java @@ -0,0 +1,40 @@ +/* + * 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.io.streaming; + +import java.io.IOException; +import java.io.UncheckedIOException; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; + +/** Creates a {@link BeamPartitionReader} for a {@link BeamInputPartition} on the executor. */ +public class BeamPartitionReaderFactory implements PartitionReaderFactory { + + private static final long serialVersionUID = 1L; + + @Override + public PartitionReader createReader(InputPartition partition) { + try { + return new BeamPartitionReader<>((BeamInputPartition) partition); + } catch (IOException e) { + throw new UncheckedIOException("Failed to open Beam reader for " + partition, e); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java new file mode 100644 index 000000000000..cbc697412b18 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -0,0 +1,277 @@ +/* + * 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.io.streaming; + +import java.io.Closeable; +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; +import org.apache.beam.sdk.options.PipelineOptions; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Executor side cache of live Beam {@link UnboundedReader}s keyed by checkpoint location and split. + * + *

An entry records the epoch its reader is positioned at and the mark taken there. A batch + * starting at that epoch reuses the reader and finalizes the pending mark, the start epoch of a + * batch is always committed by Spark. Any other start epoch, or a reader that moved without + * completing its batch, closes the entry without finalizing and restores the reader from the + * durable mark at the start epoch. + * + *

A reader idle for longer than its idle timeout is closed without finalizing its mark, the + * source redelivers. The timeout must exceed the longest gap between two micro-batches of one + * split. Speculative execution can leave a losing attempt's mark finalized on another executor, + * this source is not safe under {@code spark.speculation} with sources whose reads are not + * deterministic. + */ +public final class BeamReaderCache { + + private static final Logger LOG = LoggerFactory.getLogger(BeamReaderCache.class); + + private static final ConcurrentMap> READERS = new ConcurrentHashMap<>(); + + /** One monitor per key, acquire serializes per split, not across splits. */ + private static final ConcurrentMap LOCKS = new ConcurrentHashMap<>(); + + private BeamReaderCache() {} + + public static String key(String checkpointLocation, int splitId) { + return checkpointLocation + '|' + splitId; + } + + /** Supplies the durable coded mark at the start epoch of a batch, null if there is none. */ + @FunctionalInterface + interface MarkRestorer { + byte @Nullable [] restore() throws IOException; + } + + /** + * Returns the reader for {@code key} positioned at {@code startEpoch}, reusing the cached one if + * it is there, restoring from the durable mark otherwise. A zero length durable mark means a + * fresh start. + * + * @throws IllegalStateException if {@code startEpoch > 0} and no durable mark exists + */ + public static CachedReader acquire( + String key, + long startEpoch, + UnboundedSource source, + PipelineOptions options, + long idleTimeoutMillis, + MarkRestorer restorer) + throws IOException { + closeIdle(); + synchronized (lock(key)) { + CachedReader existing = READERS.get(key); + if (existing != null) { + if (existing.beginBatch(startEpoch)) { + existing.finalizePendingMark(key); + @SuppressWarnings("unchecked") // one source per key, its element type never changes + CachedReader reused = (CachedReader) existing; + return reused; + } + LOG.info( + "Cached Beam reader {} is at epoch {}, batch starts at {}, restoring from the durable" + + " mark.", + key, + existing.positionEpoch(), + startEpoch); + invalidate(key); + } + byte[] codedMark = restorer.restore(); + if (codedMark == null && startEpoch > 0) { + throw new IllegalStateException( + "No durable checkpoint mark for Beam reader " + key + " at epoch " + startEpoch); + } + if (codedMark != null && codedMark.length == 0) { + codedMark = null; + } + LOG.info( + "Creating Beam reader {} at epoch {} ({} mark).", + key, + startEpoch, + codedMark == null ? "no" : "restored"); + CachedReader created = + new CachedReader<>( + createReader(source, options, codedMark), startEpoch, codedMark, idleTimeoutMillis); + created.beginBatch(startEpoch); + READERS.put(key, created); + return created; + } + } + + private static UnboundedReader createReader( + UnboundedSource source, PipelineOptions options, byte @Nullable [] codedMark) + throws IOException { + MarkT mark = + codedMark == null + ? null + : CoderHelpers.fromByteArray(codedMark, source.getCheckpointMarkCoder()); + return source.createReader(options, mark); + } + + /** Closes and forgets the reader of {@code key}, nothing is finalized. */ + public static void invalidate(String key) { + synchronized (lock(key)) { + CachedReader removed = READERS.remove(key); + if (removed != null) { + close(key, removed); + } + } + } + + /** Closes and forgets every cached reader. */ + public static void invalidateAll() { + for (String key : READERS.keySet()) { + invalidate(key); + } + } + + private static void closeIdle() { + long now = System.currentTimeMillis(); + for (Map.Entry> entry : READERS.entrySet()) { + if (entry.getValue().isIdleSince(now)) { + LOG.info("Closing idle Beam reader {}.", entry.getKey()); + invalidate(entry.getKey()); + } + } + } + + private static void close(String key, CachedReader reader) { + try { + reader.close(); + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to close Beam reader {}.", key, e); + } + } + + private static Object lock(String key) { + return LOCKS.computeIfAbsent(key, k -> new Object()); + } + + /** A live reader with the epoch it is positioned at and the coded mark taken there. */ + public static final class CachedReader implements Closeable { + private final UnboundedReader reader; + private final long idleTimeoutMillis; + private boolean started; + private boolean inBatch; + private boolean moved; + private long positionEpoch; + private byte @Nullable [] positionMark; + private @Nullable CheckpointMark pendingMark; + private long lastUsedMillis; + + CachedReader( + UnboundedReader reader, + long positionEpoch, + byte @Nullable [] positionMark, + long idleTimeoutMillis) { + this.reader = reader; + this.positionEpoch = positionEpoch; + this.positionMark = positionMark; + this.idleTimeoutMillis = idleTimeoutMillis; + this.lastUsedMillis = System.currentTimeMillis(); + } + + public UnboundedReader reader() { + return reader; + } + + public synchronized boolean startOrAdvance() throws IOException { + moved = true; + if (!started) { + started = true; + return reader.start(); + } + return reader.advance(); + } + + /** + * Whether {@link #startOrAdvance()} was called at least once, only then may a mark be taken. + */ + public synchronized boolean started() { + return started; + } + + public synchronized long positionEpoch() { + return positionEpoch; + } + + /** The coded mark of the current position, null for a fresh start. */ + synchronized byte @Nullable [] positionMark() { + return positionMark; + } + + /** + * Claims the reader for a batch starting at {@code epoch}, false if it cannot continue there. + */ + synchronized boolean beginBatch(long epoch) { + if (positionEpoch != epoch || moved) { + return false; + } + inBatch = true; + lastUsedMillis = System.currentTimeMillis(); + return true; + } + + /** Records a completed batch, the reader is positioned at {@code endEpoch} from now on. */ + synchronized void endBatch(long endEpoch, @Nullable CheckpointMark mark, byte[] codedMark) { + positionEpoch = endEpoch; + positionMark = codedMark; + pendingMark = mark; + moved = false; + inBatch = false; + lastUsedMillis = System.currentTimeMillis(); + } + + synchronized boolean isIdleSince(long nowMillis) { + return !inBatch && nowMillis - lastUsedMillis > idleTimeoutMillis; + } + + /** Finalizes the pending mark if any, a failure is logged. */ + synchronized void finalizePendingMark(String key) { + CheckpointMark mark = pendingMark; + pendingMark = null; + if (mark == null) { + return; + } + LOG.debug("Finalizing checkpoint mark of Beam reader {} at epoch {}.", key, positionEpoch); + try { + mark.finalizeCheckpoint(); + } catch (IOException | RuntimeException e) { + LOG.warn( + "Failed to finalize checkpoint mark of Beam reader {} at epoch {}.", + key, + positionEpoch, + e); + } + } + + @Override + public void close() throws IOException { + reader.close(); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java new file mode 100644 index 000000000000..74b4bd101a1d --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java @@ -0,0 +1,184 @@ +/* + * 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.io.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.Path; +import org.apache.spark.sql.execution.streaming.CheckpointFileManager; +import org.apache.spark.sql.execution.streaming.CheckpointFileManager.CancellableFSDataOutputStream; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Durable state of one Beam unbounded source under the per source checkpoint location Spark hands + * to {@code toMicroBatchStream}. + * + *

{@code /splits} pins the split list, written once by the driver. {@code + * /marks//} holds the coded checkpoint mark of a split at the end of the + * batch ending at that epoch. All IO goes through Spark's {@link CheckpointFileManager}, writes are + * atomic renames. + */ +public final class BeamSourceCheckpoint { + + private static final Logger LOG = LoggerFactory.getLogger(BeamSourceCheckpoint.class); + + private static final String SPLITS_FILE = "splits"; + private static final String MARKS_DIR = "marks"; + + /** A purge further than this above the last one lists the directory instead of probing epochs. */ + private static final long MAX_BLIND_PURGE_RANGE = 1_000L; + + private final String location; + private final CheckpointFileManager fm; + private final Path root; + private final Path splitsPath; + private final Path marksRoot; + + /** Splits whose mark directory this instance already created. */ + private final Set preparedSplits = ConcurrentHashMap.newKeySet(); + + /** Per split, every mark epoch strictly below the value is known to be deleted. */ + private final ConcurrentMap purgeFloors = new ConcurrentHashMap<>(); + + public BeamSourceCheckpoint(String checkpointLocation, Configuration hadoopConf) { + this.location = checkpointLocation; + this.root = new Path(checkpointLocation); + this.fm = CheckpointFileManager.create(root, hadoopConf); + this.splitsPath = new Path(root, SPLITS_FILE); + this.marksRoot = new Path(root, MARKS_DIR); + } + + public String location() { + return location; + } + + /** The pinned split list, or null if none was pinned yet. */ + public @Nullable List> readSplits() throws IOException { + if (!fm.exists(splitsPath)) { + return null; + } + @SuppressWarnings("unchecked") // written by writeSplits as an ArrayList of sources + List> splits = + (List>) + SerializableUtils.deserializeFromByteArray(read(splitsPath), "splits at " + splitsPath); + return splits; + } + + /** Pins the split list, fails if one is pinned already. */ + public void writeSplits(List> splits) throws IOException { + fm.mkdirs(root); + if (fm.exists(splitsPath)) { + throw new IOException("Split list already pinned at " + splitsPath); + } + write(splitsPath, SerializableUtils.serializeToByteArray(new ArrayList<>(splits)), false); + LOG.info("Pinned {} split(s) at {}.", splits.size(), splitsPath); + } + + public void writeMark(int splitId, long epoch, byte[] codedMark) throws IOException { + if (preparedSplits.add(splitId)) { + fm.mkdirs(marksDir(splitId)); + } + write(markPath(splitId, epoch), codedMark, true); + } + + /** The coded mark of a split at an epoch, or null if absent. */ + public byte @Nullable [] readMark(int splitId, long epoch) throws IOException { + Path path = markPath(splitId, epoch); + if (!fm.exists(path)) { + return null; + } + return read(path); + } + + /** + * Deletes every mark of a split with an epoch strictly below {@code epoch}. Lists the directory + * once per split, later calls delete the range above the previous floor only. Idempotent. + */ + public void purgeMarksBelow(int splitId, long epoch) throws IOException { + Long floor = purgeFloors.get(splitId); + if (floor != null && epoch - floor > MAX_BLIND_PURGE_RANGE) { + floor = null; + } + if (floor == null) { + Path dir = marksDir(splitId); + if (!fm.exists(dir)) { + return; + } + for (FileStatus status : fm.list(dir)) { + long existing = parseEpoch(status.getPath().getName()); + if (existing >= 0 && existing < epoch) { + fm.delete(status.getPath()); + } + } + purgeFloors.put(splitId, epoch); + return; + } + for (long e = floor; e < epoch; e++) { + fm.delete(markPath(splitId, e)); + } + if (epoch > floor) { + purgeFloors.put(splitId, epoch); + } + } + + private Path marksDir(int splitId) { + return new Path(marksRoot, Integer.toString(splitId)); + } + + private Path markPath(int splitId, long epoch) { + return new Path(marksDir(splitId), Long.toString(epoch)); + } + + private byte[] read(Path path) throws IOException { + try (FSDataInputStream in = fm.open(path)) { + return ByteStreams.toByteArray(in); + } + } + + private void write(Path path, byte[] bytes, boolean overwrite) throws IOException { + CancellableFSDataOutputStream out = fm.createAtomic(path, overwrite); + try { + out.write(bytes); + out.close(); + } catch (IOException | RuntimeException e) { + out.cancel(); + throw e; + } + } + + /** The epoch encoded in a mark file name, or -1 for anything else. */ + private static long parseEpoch(String name) { + try { + return Long.parseLong(name); + } catch (NumberFormatException e) { + return -1L; + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.java new file mode 100644 index 000000000000..e024550d29a3 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.java @@ -0,0 +1,100 @@ +/* + * 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.io.streaming; + +import java.io.Serializable; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.util.SerializableConfiguration; + +/** Everything the driver needs to plan micro-batches of one Beam {@link UnboundedSource}. */ +final class BeamSourceSpec implements Serializable { + + private static final long serialVersionUID = 1L; + + private final UnboundedSource source; + private final Coder> coder; + private final Broadcast options; + private final Broadcast hadoopConf; + private final int desiredNumSplits; + private final long maxRecordsPerBatch; + private final long maxBatchDurationMillis; + private final long readerIdleTimeoutMillis; + private final String transformName; + + BeamSourceSpec( + UnboundedSource source, + Coder> coder, + Broadcast options, + Broadcast hadoopConf, + int desiredNumSplits, + long maxRecordsPerBatch, + long maxBatchDurationMillis, + long readerIdleTimeoutMillis, + String transformName) { + this.source = source; + this.coder = coder; + this.options = options; + this.hadoopConf = hadoopConf; + this.desiredNumSplits = desiredNumSplits; + this.maxRecordsPerBatch = maxRecordsPerBatch; + this.maxBatchDurationMillis = maxBatchDurationMillis; + this.readerIdleTimeoutMillis = readerIdleTimeoutMillis; + this.transformName = transformName; + } + + UnboundedSource source() { + return source; + } + + Coder> coder() { + return coder; + } + + Broadcast options() { + return options; + } + + Broadcast hadoopConf() { + return hadoopConf; + } + + int desiredNumSplits() { + return desiredNumSplits; + } + + /** Records per micro-batch across all splits, below 1 means unlimited. */ + long maxRecordsPerBatch() { + return maxRecordsPerBatch; + } + + long maxBatchDurationMillis() { + return maxBatchDurationMillis; + } + + long readerIdleTimeoutMillis() { + return readerIdleTimeoutMillis; + } + + String transformName() { + return transformName; + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java new file mode 100644 index 000000000000..d6fc2ada4dab --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java @@ -0,0 +1,82 @@ +/* + * 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.io.streaming; + +import java.util.Set; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.spark.sql.connector.catalog.SupportsRead; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableCapability; +import org.apache.spark.sql.connector.read.Scan; +import org.apache.spark.sql.connector.read.ScanBuilder; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; + +/** DataSourceV2 {@link Table} over a Beam unbounded source, micro-batch reads only. */ +public class BeamStreamingTable implements Table, SupportsRead { + + private final BeamSourceSpec spec; + + BeamStreamingTable(BeamSourceSpec spec) { + this.spec = spec; + } + + @Override + public String name() { + return "BeamUnboundedSource[" + spec.transformName() + "]"; + } + + @Override + public StructType schema() { + return UnboundedSourceDataset.SCHEMA; + } + + @Override + public Set capabilities() { + return ImmutableSet.of(TableCapability.MICRO_BATCH_READ); + } + + @Override + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap ignored) { + return () -> new BeamScan(spec); + } + + private static class BeamScan implements Scan { + private final BeamSourceSpec spec; + + BeamScan(BeamSourceSpec spec) { + this.spec = spec; + } + + @Override + public StructType readSchema() { + return UnboundedSourceDataset.SCHEMA; + } + + @Override + public String description() { + return "BeamUnboundedSource[" + spec.transformName() + "]"; + } + + @Override + public MicroBatchStream toMicroBatchStream(String checkpointLocation) { + return new BeamMicroBatchStream<>(spec, checkpointLocation); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java new file mode 100644 index 000000000000..87a0fa0aeff6 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -0,0 +1,121 @@ +/* + * 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.io.streaming; + +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan; +import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2; +import org.apache.spark.sql.catalyst.types.DataTypeUtils; +import org.apache.spark.sql.classic.Dataset$; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.apache.spark.util.SerializableConfiguration; +import scala.Option; +import scala.reflect.ClassTag; + +/** + * Translator facing entry point turning a Beam {@link UnboundedSource} into a streaming Spark + * {@link Dataset} of rows. + * + *

The dataset has two columns, {@value #COL_PAYLOAD} of type {@code BINARY} holding the element + * encoded with the supplied {@code WindowedValue} coder, and {@value #COL_EVENT_TS} of type {@code + * TIMESTAMP} holding the event timestamp of that element. + * + *

The event time watermark is declared here and only here. Spark 4 rejects a second {@code + * withWatermark} further down the plan, so downstream translators must never call it again. + */ +public final class UnboundedSourceDataset { + + public static final String COL_PAYLOAD = "payload"; + + public static final String COL_EVENT_TS = "eventTimestamp"; + + public static final StructType SCHEMA = + new StructType() + .add(COL_PAYLOAD, DataTypes.BinaryType, false) + .add(COL_EVENT_TS, DataTypes.TimestampType, false); + + private static final String SOURCE_NAME = "beam-unbounded"; + + private UnboundedSourceDataset() {} + + /** + * Builds the streaming {@link Dataset} for {@code source} with the event time watermark applied. + * + * @param session the active Spark session + * @param source the Beam unbounded source to read + * @param windowedValueCoder the coder of the {@value #COL_PAYLOAD} column + * @param options the pipeline options, supplying the watermark delay and the micro-batch limits + * @param transformName the full name of the read transform, used for naming only + * @param the element type of the source + * @param the checkpoint mark type of the source + */ + public static Dataset of( + SparkSession session, + UnboundedSource source, + Coder> windowedValueCoder, + SparkStructuredStreamingPipelineOptions options, + String transformName) { + org.apache.spark.sql.classic.SparkSession classic = + (org.apache.spark.sql.classic.SparkSession) session; + Configuration hadoopConf = classic.sessionState().newHadoopConf(); + BeamSourceSpec spec = + new BeamSourceSpec<>( + source, + windowedValueCoder, + broadcast( + session, + new SerializablePipelineOptions(options), + SerializablePipelineOptions.class), + broadcast( + session, + new SerializableConfiguration(hadoopConf), + SerializableConfiguration.class), + session.sparkContext().defaultParallelism(), + options.getMaxRecordsPerBatch(), + Math.max(1L, options.getMaxBatchDurationMillis()), + options.getReaderIdleTimeoutMillis(), + transformName); + LogicalPlan plan = + new StreamingRelationV2( + Option.empty(), + SOURCE_NAME, + new BeamStreamingTable(spec), + CaseInsensitiveStringMap.empty(), + DataTypeUtils.toAttributes(SCHEMA), + Option.empty(), + Option.empty(), + Option.empty()); + Dataset rows = Dataset$.MODULE$.ofRows(classic, plan); + return rows.withWatermark(COL_EVENT_TS, options.getWatermarkDelayMillis() + " milliseconds"); + } + + private static Broadcast broadcast(SparkSession session, T value, Class type) { + return session.sparkContext().broadcast(value, ClassTag.apply(type)); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java new file mode 100644 index 000000000000..cc7007267601 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -0,0 +1,1146 @@ +/* + * 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.io.streaming; + +import static org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.COL_EVENT_TS; +import static org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.COL_PAYLOAD; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.Serializable; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.sql.Timestamp; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; +import javax.annotation.Nullable; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +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.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.api.java.function.VoidFunction2; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark; +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryProgress; +import org.apache.spark.sql.streaming.Trigger; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.apache.spark.util.SerializableConfiguration; +import org.joda.time.Instant; +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; +import scala.reflect.ClassTag; + +/** + * Tests for the Spark 4 DataSourceV2 micro-batch source wrapping a Beam {@link UnboundedSource}. + * + *

Termination note: the epoch offsets of this source never settle, {@code + * StreamingQuery.processAllAvailable()} would therefore block forever. Every test here drives the + * query with {@code Trigger.ProcessingTime(100)} and stops it explicitly once the expected result + * arrived or the poll deadline expired. + */ +@Category(StreamingTest.class) +@RunWith(JUnit4.class) +public class BeamMicroBatchSourceTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder temp = new TemporaryFolder(); + + private static final AtomicInteger QUERY_COUNTER = new AtomicInteger(); + + /** Rows collected per query name by {@link #startCollecting}, driver side. */ + private static final Map> COLLECTED = new ConcurrentHashMap<>(); + + /** Rows per micro-batch per query name, see {@link #startCollectingBatches}, driver side. */ + private static final Map>> BATCHES = new ConcurrentHashMap<>(); + + /** 2023-11-14T22:13:20Z, a plain modern timestamp with no rebase or DST subtleties. */ + private static final long BASE_MILLIS = 1_700_000_000_000L; + + private static final long INTERVAL_MILLIS = 1_000L; + + private static final long POLL_TIMEOUT_MILLIS = 120_000L; + + @After + public void tearDown() { + BeamReaderCache.invalidateAll(); + COLLECTED.clear(); + BATCHES.clear(); + ShardedListSource.FINALIZED.clear(); + } + + /** + * THE TOP RISK OF THE POC: does Spark's {@code EventTimeWatermark} logical node survive a typed + * transformation producing a Beam typed dataset, so a later stateful operator can still see it? + * + *

Verdict, observed on Spark 4: it does. The node stays at the bottom of both the logical and + * the analyzed plan through one and through two typed maps, {@code DeserializeToObject / + * MapElements / SerializeFromObject} are simply stacked on top of it. + * + *

One nuance the translators must be aware of: the per attribute delay marker (rendered as + * {@code eventTimestamp#1-T1000ms} in the analyzed plan) only travels with the timestamp + * attribute itself. Once a typed map projects the timestamp column away, downstream attributes + * carry no marker. Operators that read the marker off an attribute, for example {@code + * groupBy(window(...))} or a stream to stream join, would therefore not see it. Operators that + * read the query wide watermark, which is what {@code transformWithState} in {@code + * TimeMode.EventTime} does, are unaffected, see {@link + * #testWatermarkIsTrackedAtRuntimeAfterTypedMap}. + */ + @Test + public void testEventTimeWatermarkSurvivesTypedMap() { + Dataset rows = rows(4, 1_000L); + assertTrue("source dataset must be streaming", rows.isStreaming()); + + assertWatermark("directly after withWatermark, logical plan", logical(rows)); + assertWatermark("directly after withWatermark, analyzed plan", analyzed(rows)); + + // A typed map producing a Beam-ish (opaque bytes) dataset, exactly what the read translator + // will do to turn rows into WindowedValue bytes. + Dataset typed = + rows.map((MapFunction) row -> row.getAs(COL_PAYLOAD), Encoders.BINARY()); + + assertWatermark("after a typed map, logical plan", logical(typed)); + assertWatermark("after a typed map, analyzed plan", analyzed(typed)); + + // And once more after a second typed stage, mimicking chained operators. + Dataset chained = + typed.map((MapFunction) bytes -> bytes, Encoders.BINARY()); + + assertWatermark("after two chained typed maps, logical plan", logical(chained)); + assertWatermark("after two chained typed maps, analyzed plan", analyzed(chained)); + } + + /** + * Runtime counterpart of the plan inspection above, the watermark must actually be tracked by the + * running query, not merely present as a plan node. + */ + @Test + public void testWatermarkIsTrackedAtRuntimeAfterTypedMap() throws Exception { + Dataset rows = rows(8, 0L); + Dataset typed = + rows.map((MapFunction) row -> row.getAs(COL_PAYLOAD), Encoders.BINARY()); + + String queryName = "beam_wm_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = startDiscarding(typed, queryName); + try { + String watermark = awaitWatermark(query); + assertNotNull("query never reported an event time watermark", watermark); + assertFalse( + "watermark never advanced past the epoch, it is not being tracked: " + watermark, + watermark.startsWith("1970-")); + } finally { + stopQuietly(query); + } + } + + /** Reads a finite set of elements through the DSv2 source and checks payloads and timestamps. */ + @Test + public void testReadsElementsFromUnboundedSource() throws Exception { + int count = 8; + Dataset rows = rows(count, 0L); + + String queryName = "beam_read_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = startCollecting(rows, queryName); + List collected; + try { + collected = awaitRows(queryName, count); + } finally { + stopQuietly(query); + } + + assertEquals("unexpected number of rows", count, collected.size()); + + Coder> coder = coder(); + List values = new ArrayList<>(); + for (Row row : collected) { + byte[] payload = row.getAs(COL_PAYLOAD); + Timestamp eventTs = row.getAs(COL_EVENT_TS); + WindowedValue windowedValue = CoderUtils.decodeFromByteArray(coder, payload); + values.add(windowedValue.getValue()); + assertEquals( + "eventTimestamp column must match the timestamp inside the encoded WindowedValue", + windowedValue.getTimestamp().getMillis(), + eventTs.getTime()); + assertEquals( + "elements are read into the global window", + Collections.singletonList(GlobalWindow.INSTANCE), + new ArrayList<>(windowedValue.getWindows())); + } + + Collections.sort(values); + List expected = new ArrayList<>(); + for (int i = 0; i < count; i++) { + expected.add(element(i)); + } + assertEquals(expected, values); + } + + /** The default record limit of -1 means unlimited, an available source drains in one batch. */ + @Test + public void testUnlimitedRecordsPerBatchByDefault() throws Exception { + int count = 2500; + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(0L); + options.setMaxBatchDurationMillis(5_000L); + Dataset rows = + UnboundedSourceDataset.of( + SESSION.getSession(), new ListSource(count), coder(), options, "Read(ListSource)"); + + String queryName = "beam_nolimit_" + QUERY_COUNTER.incrementAndGet(); + List batchSizes = Collections.synchronizedList(new ArrayList<>()); + StreamingQuery query = + rows.writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + long size = batch.count(); + if (size > 0) { + batchSizes.add(size); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + try { + long deadline = System.currentTimeMillis() + POLL_TIMEOUT_MILLIS; + while (System.currentTimeMillis() < deadline && batchSizes.isEmpty()) { + Thread.sleep(100L); + } + } finally { + stopQuietly(query); + } + assertEquals( + "without a limit all elements must arrive in the first non-empty micro-batch", + Collections.singletonList((long) count), + new ArrayList<>(batchSizes)); + } + + /** The offset is an opaque, strictly increasing epoch counter, its JSON is the bare number. */ + @Test + public void testEpochOffsetRoundTrip() { + BeamOffset offset = new BeamOffset(42L); + assertEquals("42", offset.json()); + assertEquals(42L, BeamOffset.fromJson("42").epoch()); + assertEquals(0L, BeamOffset.ZERO.epoch()); + assertEquals(new BeamOffset(7L), new BeamOffset(7L)); + assertThrows(IllegalArgumentException.class, () -> BeamOffset.fromJson("x")); + } + + @Test + public void testEpochFastForwardsPastDeserializedOffset() throws Exception { + BeamMicroBatchStream stream = newStream(temp.newFolder("ff-offset").getAbsolutePath()); + stream.deserializeOffset("7"); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the replayed epoch 7, got " + next, next.epoch() > 7L); + } + + @Test + public void testEpochFastForwardsPastPlannedOffsets() throws Exception { + BeamMicroBatchStream stream = newStream(temp.newFolder("ff-plan").getAbsolutePath()); + InputPartition[] partitions = + stream.planInputPartitions(new BeamOffset(3L), new BeamOffset(9L)); + assertTrue("at least one partition expected", partitions.length > 0); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the planned epoch 9, got " + next, next.epoch() > 9L); + } + + /** The batch quota is divided over the splits, remainder first, every split gets at least one. */ + @Test + public void testSplitQuotas() { + assertArrayEquals(new long[] {1, 1, 1, 1, 1, 1, 1, 1}, BeamMicroBatchStream.splitQuotas(3, 8)); + assertArrayEquals(new long[] {4, 3, 3}, BeamMicroBatchStream.splitQuotas(10, 3)); + assertArrayEquals(new long[] {5, 5}, BeamMicroBatchStream.splitQuotas(10, 2)); + assertArrayEquals(new long[] {0, 0, 0}, BeamMicroBatchStream.splitQuotas(0, 3)); + assertArrayEquals(new long[] {-1, -1, -1}, BeamMicroBatchStream.splitQuotas(-1, 3)); + long[] many = BeamMicroBatchStream.splitQuotas(1, 200); + assertEquals(200, many.length); + for (long quota : many) { + assertEquals(1L, quota); + } + } + + /** Rendezvous hashing: a joining executor only takes splits for itself, order does not matter. */ + @Test + public void testSplitAssignmentIsStableWhenExecutorJoins() { + List three = Arrays.asList("executor_a_1", "executor_b_2", "executor_c_3"); + List shuffled = Arrays.asList("executor_c_3", "executor_a_1", "executor_b_2"); + List four = new ArrayList<>(three); + four.add("executor_d_4"); + int splits = 200; + int moved = 0; + for (int split = 0; split < splits; split++) { + String before = BeamMicroBatchStream.assign(split, three); + String after = BeamMicroBatchStream.assign(split, four); + assertTrue(three.contains(before)); + assertEquals(before, BeamMicroBatchStream.assign(split, shuffled)); + if (!after.equals(before)) { + assertEquals("split " + split + " moved to an old executor", "executor_d_4", after); + moved++; + } + } + assertTrue("the new executor took no split", moved > 0); + assertTrue("the new executor took every split", moved < splits); + assertEquals( + "executor_a_1", BeamMicroBatchStream.assign(7, Collections.singletonList("executor_a_1"))); + } + + /** The record limit of a micro-batch is a total over all splits, not a per split allowance. */ + @Test + public void testMaxRecordsPerBatchIsSharedAcrossSplits() throws Exception { + int shards = 2; + int count = 30; + long limit = 10L; + String queryName = "beam_shared_limit_" + QUERY_COUNTER.incrementAndGet(); + Dataset rows = shardedRows(queryName, shards, count, limit); + StreamingQuery query = startCollectingBatches(rows, queryName, temp.newFolder(queryName)); + try { + await("all rows", () -> values(batches(queryName)).size() >= count); + } finally { + stopQuietly(query); + } + + List sizes = new ArrayList<>(); + for (List batch : batches(queryName)) { + if (!batch.isEmpty()) { + sizes.add(batch.size()); + } + } + assertFalse("no rows arrived", sizes.isEmpty()); + assertTrue("first batch exceeds the shared limit: " + sizes, sizes.get(0) <= limit); + for (int size : sizes) { + assertTrue("batch exceeds the shared limit: " + sizes, size <= limit); + } + List values = values(batches(queryName)); + assertEquals(count, values.size()); + assertEquals(ShardedListSource.elements(shards, count), new HashSet<>(values)); + } + + /** + * A restart resumes every split from the durable mark of the last committed batch, replaying at + * most the one uncommitted batch. + */ + @Test + public void testRestartResumesFromCommittedMark() throws Exception { + int shards = 2; + int count = 80; + long limit = 4L; + File checkpointDir = temp.newFolder("restart"); + Set all = ShardedListSource.elements(shards, count); + + String first = "beam_restart_a_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = + startCollectingBatches(shardedRows(first, shards, count, limit), first, checkpointDir); + try { + await("two commits", () -> committedBatchIds(checkpointDir).size() >= 2); + } finally { + stopQuietly(query); + } + BeamReaderCache.invalidateAll(); + List firstValues = values(batches(first)); + + String second = "beam_restart_b_" + QUERY_COUNTER.incrementAndGet(); + query = + startCollectingBatches(shardedRows(second, shards, count, limit), second, checkpointDir); + try { + await( + "union of both runs", + () -> { + Set union = new HashSet<>(firstValues); + union.addAll(values(batches(second))); + return union.containsAll(all); + }); + } finally { + stopQuietly(query); + } + List secondValues = values(batches(second)); + + Set union = new HashSet<>(firstValues); + union.addAll(secondValues); + assertEquals(all, union); + assertTrue( + "more than the uncommitted batch replayed: " + + firstValues.size() + + " + " + + secondValues.size() + + " rows", + firstValues.size() + secondValues.size() <= count + limit); + for (int shard = 0; shard < shards; shard++) { + int min = Integer.MAX_VALUE; + for (String value : secondValues) { + if (ShardedListSource.shardOf(value) == shard) { + min = Math.min(min, ShardedListSource.indexOf(value)); + } + } + assertTrue("run 2 delivered nothing for shard " + shard, min < Integer.MAX_VALUE); + assertTrue("run 2 restarted shard " + shard + " from element 0", min > 0); + } + } + + /** + * Every finalized position of a split is at most the position in that split's mark at the end + * epoch of the highest committed batch, so no mark is finalized before Spark commits its batch. + */ + @Test + public void testMarksAreFinalizedOnlyAfterSparkCommit() throws Exception { + int shards = 2; + // Never exhausted while the test runs, positions strictly increase with the epoch. + int count = 4_000; + File checkpointDir = temp.newFolder("finalize"); + String queryName = "beam_finalize_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = + startCollectingBatches(shardedRows(queryName, shards, count, 4L), queryName, checkpointDir); + try { + await("three commits", () -> committedBatchIds(checkpointDir).size() >= 3); + } finally { + stopQuietly(query); + } + + long committedEpoch = + endEpoch(checkpointDir, Collections.max(committedBatchIds(checkpointDir))); + BeamSourceCheckpoint files = sourceCheckpoint(checkpointDir); + int finalizations = 0; + for (int shard = 0; shard < shards; shard++) { + byte[] coded = files.readMark(shard, committedEpoch); + assertNotNull("no mark at committed epoch " + committedEpoch + " for split " + shard, coded); + int committedPosition = + CoderUtils.decodeFromByteArray(ShardedListSource.MARK_CODER, coded).next; + List finalized = ShardedListSource.finalized(queryName, shard); + for (int position : finalized) { + assertTrue( + "split " + + shard + + " finalized position " + + position + + " beyond committed position " + + committedPosition, + position <= committedPosition); + } + finalizations += finalized.size(); + } + assertTrue("no mark was finalized", finalizations > 0); + } + + /** + * Spark reports the commit of batch N-1 when it constructs batch N, so after a run the surviving + * marks of every split are at or above the end epoch of the batch before the last constructed + * one, the mark at the highest committed end epoch exists, and the mark of batch 0 is gone. + */ + @Test + public void testMarksBelowCommittedOffsetArePurged() throws Exception { + int shards = 2; + int count = 4_000; + File checkpointDir = temp.newFolder("purge"); + String queryName = "beam_purge_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = + startCollectingBatches(shardedRows(queryName, shards, count, 4L), queryName, checkpointDir); + try { + await("three commits", () -> committedBatchIds(checkpointDir).size() >= 3); + } finally { + stopQuietly(query); + } + + long lastConstructed = Collections.max(batchIds(new File(checkpointDir, "offsets"))); + long purgeFloor = endEpoch(checkpointDir, lastConstructed - 1); + long committedEpoch = + endEpoch(checkpointDir, Collections.max(committedBatchIds(checkpointDir))); + long firstEpoch = endEpoch(checkpointDir, 0); + assertTrue(committedEpoch >= purgeFloor); + assertTrue(purgeFloor > firstEpoch); + File sourceDir = new File(checkpointDir, "sources/0"); + // The last purge runs asynchronously on the driver and may still be in flight after stop. + awaitQuietly( + 10_000L, + () -> { + for (int shard = 0; shard < shards; shard++) { + for (long epoch : batchIds(new File(sourceDir, "marks/" + shard))) { + if (epoch < purgeFloor) { + return false; + } + } + } + return true; + }); + for (int shard = 0; shard < shards; shard++) { + Set remaining = batchIds(new File(sourceDir, "marks/" + shard)); + assertTrue( + "split " + + shard + + " lost the mark at committed epoch " + + committedEpoch + + ": " + + remaining, + remaining.contains(committedEpoch)); + for (long epoch : remaining) { + assertTrue( + "split " + shard + " kept mark " + epoch + " below purge floor " + purgeFloor, + epoch >= purgeFloor); + } + assertFalse("split " + shard + " kept the mark of batch 0", remaining.contains(firstEpoch)); + } + } + + // --------------------------------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------------------------------- + + private Dataset rows(int count, long watermarkDelayMillis) { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(watermarkDelayMillis); + options.setMaxRecordsPerBatch(1000L); + options.setMaxBatchDurationMillis(200L); + return UnboundedSourceDataset.of( + SESSION.getSession(), new ListSource(count), coder(), options, "Read(ListSource)"); + } + + /** Builds the driver side stream through the table, with real broadcasts from the session. */ + private static BeamMicroBatchStream newStream(String checkpointLocation) { + SparkSession session = SESSION.getSession(); + Broadcast options = + session + .sparkContext() + .broadcast( + new SerializablePipelineOptions(PipelineOptionsFactory.create()), + ClassTag.apply(SerializablePipelineOptions.class)); + Broadcast hadoopConf = + session + .sparkContext() + .broadcast( + new SerializableConfiguration(new Configuration()), + ClassTag.apply(SerializableConfiguration.class)); + BeamSourceSpec spec = + new BeamSourceSpec<>( + new ListSource(4), + coder(), + options, + hadoopConf, + 2, + -1L, + 200L, + 600_000L, + "Read(ListSource)"); + MicroBatchStream stream = + new BeamStreamingTable(spec) + .newScanBuilder(CaseInsensitiveStringMap.empty()) + .build() + .toMicroBatchStream(checkpointLocation); + return (BeamMicroBatchStream) stream; + } + + private Dataset shardedRows(String tag, int shards, int count, long maxRecordsPerBatch) { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(0L); + options.setMaxRecordsPerBatch(maxRecordsPerBatch); + options.setMaxBatchDurationMillis(1_000L); + return UnboundedSourceDataset.of( + SESSION.getSession(), + new ShardedListSource(tag, shards, count), + coder(), + options, + "Read(ShardedListSource)"); + } + + /** Starts a query collecting every micro-batch as one list into {@link #BATCHES}. */ + private static StreamingQuery startCollectingBatches( + Dataset dataset, String queryName, File checkpointDir) throws Exception { + BATCHES.put(queryName, Collections.synchronizedList(new ArrayList<>())); + return dataset + .writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + List> target = BATCHES.get(queryName); + if (target != null) { + target.add(batch.collectAsList()); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", checkpointDir.getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + } + + private static List> batches(String queryName) { + List> batches = BATCHES.getOrDefault(queryName, Collections.emptyList()); + synchronized (batches) { + return new ArrayList<>(batches); + } + } + + private static List values(List> batches) { + Coder> coder = coder(); + List values = new ArrayList<>(); + for (List batch : batches) { + for (Row row : batch) { + byte[] payload = row.getAs(COL_PAYLOAD); + try { + values.add(CoderUtils.decodeFromByteArray(coder, payload).getValue()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + } + return values; + } + + private static void await(String what, BooleanSupplier condition) throws Exception { + if (!awaitQuietly(POLL_TIMEOUT_MILLIS, condition)) { + throw new AssertionError("timed out waiting for " + what); + } + } + + private static boolean awaitQuietly(long timeoutMillis, BooleanSupplier condition) + throws Exception { + long deadline = System.currentTimeMillis() + timeoutMillis; + while (System.currentTimeMillis() < deadline) { + if (condition.getAsBoolean()) { + return true; + } + Thread.sleep(50L); + } + return condition.getAsBoolean(); + } + + /** Numeric file names in a Spark log directory, temp and hidden files excluded. */ + private static Set batchIds(File dir) { + Set ids = new TreeSet<>(); + String[] names = dir.list(); + if (names == null) { + return ids; + } + for (String name : names) { + if (!name.startsWith(".") && !name.endsWith(".tmp")) { + try { + ids.add(Long.parseLong(name)); + } catch (NumberFormatException e) { + // not a log entry + } + } + } + return ids; + } + + private static Set committedBatchIds(File checkpointDir) { + return batchIds(new File(checkpointDir, "commits")); + } + + /** The end epoch of a batch, the offset line of the single source in {@code offsets/}. */ + private static long endEpoch(File checkpointDir, long batchId) throws IOException { + File file = new File(new File(checkpointDir, "offsets"), Long.toString(batchId)); + List lines = new ArrayList<>(); + for (String line : Files.readAllLines(file.toPath(), StandardCharsets.UTF_8)) { + if (!line.trim().isEmpty()) { + lines.add(line.trim()); + } + } + assertTrue("offset log entry too short: " + lines, lines.size() >= 3); + assertEquals("one source expected in " + lines, 3, lines.size()); + return BeamOffset.fromJson(lines.get(2)).epoch(); + } + + private static BeamSourceCheckpoint sourceCheckpoint(File checkpointDir) { + return new BeamSourceCheckpoint( + new File(checkpointDir, "sources/0").getAbsolutePath(), new Configuration()); + } + + private static Coder> coder() { + return WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + } + + private static String element(int index) { + return "element-" + index; + } + + private static long timestampMillis(int index) { + return BASE_MILLIS + index * INTERVAL_MILLIS; + } + + /** + * Starts a query that throws its output away. + * + *

The {@code noop} sink is used on purpose: these tests observe the source through a {@code + * foreachBatch} or through the query's own progress, never through the sink, so there is no + * reason to buffer rows anywhere. That matches what the streaming evaluation context does for + * real pipelines. + */ + private StreamingQuery startDiscarding(Dataset dataset, String queryName) throws Exception { + return dataset + .writeStream() + .format("noop") + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + } + + /** + * Starts a query collecting every micro-batch into {@link #COLLECTED} under {@code queryName}. + */ + private StreamingQuery startCollecting(Dataset dataset, String queryName) throws Exception { + COLLECTED.put(queryName, Collections.synchronizedList(new ArrayList<>())); + return dataset + .writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + // Exactly one action per micro-batch: any second action would re-execute the + // batch and advance the Beam reader past records that were never collected. + List target = COLLECTED.get(queryName); + if (target != null) { + target.addAll(batch.collectAsList()); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + } + + private static void stopQuietly(StreamingQuery query) { + try { + query.stop(); + } catch (Exception e) { + // Nothing useful to do while tearing a test query down. + } + } + + /** Polls the collected batches until {@code expected} rows arrived or the deadline expires. */ + private static List awaitRows(String queryName, int expected) throws Exception { + long deadline = System.currentTimeMillis() + POLL_TIMEOUT_MILLIS; + List rows = COLLECTED.getOrDefault(queryName, Collections.emptyList()); + while (System.currentTimeMillis() < deadline) { + synchronized (rows) { + if (rows.size() >= expected) { + return new ArrayList<>(rows); + } + } + Thread.sleep(100L); + } + synchronized (rows) { + return new ArrayList<>(rows); + } + } + + /** Polls the query progress until it reports an event time watermark past the epoch. */ + private static @Nullable String awaitWatermark(StreamingQuery query) throws Exception { + long deadline = System.currentTimeMillis() + POLL_TIMEOUT_MILLIS; + String last = null; + while (System.currentTimeMillis() < deadline) { + for (StreamingQueryProgress progress : query.recentProgress()) { + Map eventTime = progress.eventTime(); + String watermark = eventTime.get("watermark"); + if (watermark != null) { + last = watermark; + if (!watermark.startsWith("1970-")) { + return watermark; + } + } + } + Thread.sleep(100L); + } + return last; + } + + private static LogicalPlan logical(Dataset dataset) { + return ((org.apache.spark.sql.classic.Dataset) dataset).queryExecution().logical(); + } + + private static LogicalPlan analyzed(Dataset dataset) { + return ((org.apache.spark.sql.classic.Dataset) dataset).queryExecution().analyzed(); + } + + private static void assertWatermark(String what, LogicalPlan plan) { + assertTrue( + "no EventTimeWatermark node found " + what + ":\n" + plan.treeString(), + containsWatermark(plan)); + } + + private static boolean containsWatermark(LogicalPlan plan) { + if (plan instanceof EventTimeWatermark) { + return true; + } + scala.collection.Iterator children = plan.children().iterator(); + while (children.hasNext()) { + if (containsWatermark(children.next())) { + return true; + } + } + return false; + } + + // --------------------------------------------------------------------------------------------- + // a minimal in-memory UnboundedSource + // --------------------------------------------------------------------------------------------- + + /** + * A trivial single split {@link UnboundedSource} over a fixed number of synthetic elements with + * evenly spaced event timestamps. It is exhausted after the last element, further calls to {@code + * advance()} simply report that no data is available. + */ + private static class ListSource extends UnboundedSource { + private static final long serialVersionUID = 1L; + + private final int count; + + ListSource(int count) { + this.count = count; + } + + @Override + public List split(int desiredNumSplits, PipelineOptions options) { + return Arrays.asList(this); + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark mark) { + return new ListReader(this, mark == null ? 0 : mark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return SerializableCoder.of(Mark.class); + } + + @Override + public Coder getOutputCoder() { + return StringUtf8Coder.of(); + } + + /** Position of the next element to read. */ + static class Mark implements UnboundedSource.CheckpointMark, Serializable { + private static final long serialVersionUID = 1L; + + private final int next; + + Mark(int next) { + this.next = next; + } + + @Override + public void finalizeCheckpoint() {} + } + + private static class ListReader extends UnboundedReader { + private final ListSource source; + private int next; + private int current = -1; + + ListReader(ListSource source, int next) { + this.source = source; + this.next = next; + } + + @Override + public boolean start() { + return advance(); + } + + @Override + public boolean advance() { + if (next < source.count) { + current = next++; + return true; + } + return false; + } + + @Override + public String getCurrent() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return element(current); + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant(timestampMillis(current)); + } + + @Override + public Instant getWatermark() { + return current < 0 + ? BoundedWindow.TIMESTAMP_MIN_VALUE + : new Instant(timestampMillis(current)); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new Mark(next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() throws IOException {} + } + } + + // --------------------------------------------------------------------------------------------- + // a multi split in-memory UnboundedSource with finalization counting marks + // --------------------------------------------------------------------------------------------- + + /** + * Splits into one sub source per shard, each over {@code count / shards} elements named {@code + * shard-N-element-M} with evenly spaced timestamps. Marks are not Java serializable, they carry + * the shard's read position and record it under {@code /} when finalized. + */ + static class ShardedListSource extends UnboundedSource { + private static final long serialVersionUID = 1L; + + static final Coder MARK_CODER = new MarkCoder(); + + /** Finalized positions keyed by {@code /}. */ + static final ConcurrentMap> FINALIZED = new ConcurrentHashMap<>(); + + private final String tag; + private final int shard; + private final int shards; + private final int perShard; + + ShardedListSource(String tag, int shards, int count) { + this(tag, -1, shards, count / shards); + } + + private ShardedListSource(String tag, int shard, int shards, int perShard) { + this.tag = tag; + this.shard = shard; + this.shards = shards; + this.perShard = perShard; + } + + static Set elements(int shards, int count) { + Set elements = new HashSet<>(); + for (int shard = 0; shard < shards; shard++) { + for (int index = 0; index < count / shards; index++) { + elements.add(element(shard, index)); + } + } + return elements; + } + + static String element(int shard, int index) { + return "shard-" + shard + "-element-" + index; + } + + static int shardOf(String element) { + return Integer.parseInt(element.substring("shard-".length(), element.indexOf("-element-"))); + } + + static int indexOf(String element) { + return Integer.parseInt(element.substring(element.lastIndexOf('-') + 1)); + } + + static List finalized(String tag, int shard) { + List positions = FINALIZED.get(key(tag, shard)); + if (positions == null) { + return Collections.emptyList(); + } + synchronized (positions) { + return new ArrayList<>(positions); + } + } + + private static String key(String tag, int shard) { + return tag + "/" + shard; + } + + @Override + public List split(int desiredNumSplits, PipelineOptions options) { + if (shard >= 0) { + return Collections.singletonList(this); + } + List splits = new ArrayList<>(); + for (int i = 0; i < shards; i++) { + splits.add(new ShardedListSource(tag, i, shards, perShard)); + } + return splits; + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable ShardMark mark) { + if (shard < 0) { + throw new IllegalStateException("split before reading"); + } + return new ShardReader(this, mark == null ? 0 : mark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return MARK_CODER; + } + + @Override + public Coder getOutputCoder() { + return StringUtf8Coder.of(); + } + + /** Position of the next element of a shard, deliberately not {@link Serializable}. */ + static final class ShardMark implements UnboundedSource.CheckpointMark { + private final String key; + final int next; + + ShardMark(String key, int next) { + this.key = key; + this.next = next; + } + + @Override + public void finalizeCheckpoint() { + FINALIZED + .computeIfAbsent(key, k -> Collections.synchronizedList(new ArrayList<>())) + .add(next); + } + } + + private static final class MarkCoder extends CustomCoder { + private static final long serialVersionUID = 1L; + + @Override + public void encode(ShardMark mark, OutputStream out) throws IOException { + StringUtf8Coder.of().encode(mark.key, out); + VarIntCoder.of().encode(mark.next, out); + } + + @Override + public ShardMark decode(InputStream in) throws IOException { + return new ShardMark(StringUtf8Coder.of().decode(in), VarIntCoder.of().decode(in)); + } + } + + private static final class ShardReader extends UnboundedReader { + private final ShardedListSource source; + private int next; + private int current = -1; + + ShardReader(ShardedListSource source, int next) { + this.source = source; + this.next = next; + } + + @Override + public boolean start() { + return advance(); + } + + @Override + public boolean advance() { + if (next < source.perShard) { + current = next++; + return true; + } + return false; + } + + @Override + public String getCurrent() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return element(source.shard, current); + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant(timestampMillis(source.shard * source.perShard + current)); + } + + @Override + public Instant getWatermark() { + return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE : getCurrentTimestamp(); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new ShardMark(key(source.tag, source.shard), next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() {} + } + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java new file mode 100644 index 000000000000..138d44acdab8 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java @@ -0,0 +1,478 @@ +/* + * 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.io.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.util.SerializableConfiguration; +import org.joda.time.Instant; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +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; +import scala.reflect.ClassTag; + +/** + * Drives {@link BeamPartitionReader} directly over hand built partitions to prove the reader cache + * protocol. The session only supplies the broadcasts, no query runs. + */ +@Category(StreamingTest.class) +@RunWith(JUnit4.class) +public class BeamReaderCacheProtocolTest { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + private static final AtomicInteger TAGS = new AtomicInteger(); + + private static final Coder> CODER = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + + private static final long MAX_RECORDS = 3L; + + private static final long MAX_BATCH_MILLIS = 30_000L; + + private static final long READER_IDLE_MILLIS = 600_000L; + + private static Broadcast options; + private static Broadcast hadoopConf; + private static Configuration conf; + + private String tag; + private IntListSource source; + private String location; + + @BeforeClass + public static void broadcastOnce() { + SparkSession session = SESSION.getSession(); + conf = ((org.apache.spark.sql.classic.SparkSession) session).sessionState().newHadoopConf(); + options = + session + .sparkContext() + .broadcast( + new SerializablePipelineOptions(PipelineOptionsFactory.create()), + ClassTag.apply(SerializablePipelineOptions.class)); + hadoopConf = + session + .sparkContext() + .broadcast( + new SerializableConfiguration(conf), + ClassTag.apply(SerializableConfiguration.class)); + } + + @Before + public void setUp() throws IOException { + tag = "protocol-" + TAGS.incrementAndGet(); + source = new IntListSource(tag, 100); + location = temp.newFolder("source").getAbsolutePath(); + } + + @After + public void tearDown() { + BeamReaderCache.invalidateAll(); + IntListSource.forget(tag); + } + + /** + * A batch starting where the cached reader stopped reuses it and finalizes the mark taken there. + */ + @Test + public void testContinuationReusesReaderAndFinalizesOnce() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals("nothing is finalized before the next batch opens", noPositions(), finalized()); + + BeamPartitionReader second = open(1, 2); + assertEquals("the epoch 1 mark is finalized on open", positions(3), finalized()); + assertEquals(Arrays.asList(3, 4, 5), drain(second)); + assertEquals(positions(3), finalized()); + assertEquals("one reader serves both batches", 1, IntListSource.created(tag)); + } + + /** A retried batch restarts from the durable mark at its start and finalizes nothing. */ + @Test + public void testRetryRecreatesReaderWithoutFinalizing() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals(noPositions(), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + /** After the cache is lost the durable mark at the start epoch restores the position. */ + @Test + public void testExecutorHopRestoresFromDurableMark() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals(Arrays.asList(3, 4, 5), readBatch(1, 2)); + BeamReaderCache.invalidateAll(); + + assertEquals(Arrays.asList(6, 7, 8), readBatch(2, 3)); + assertEquals( + "only the epoch 1 mark was live when its batch was committed", positions(3), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + /** + * A start epoch above zero without a durable mark is an invariant violation, not a fresh start. + */ + @Test + public void testMissingMarkThrows() { + assertThrows(IllegalStateException.class, () -> new BeamPartitionReader<>(partition(5, 6))); + assertThrows( + IllegalStateException.class, + () -> new BeamPartitionReaderFactory().createReader(partition(5, 6))); + assertEquals(0, IntListSource.created(tag)); + } + + /** + * A reader closed before it started writes its start mark forward at E and stays reusable from E, + * a later batch from S is a fresh start. + */ + @Test + public void testNeverStartedReaderWritesStartMarkForward() throws Exception { + BeamPartitionReader idle = open(0, 1); + idle.close(); + assertTrue(new File(location, "marks/0/1").exists()); + assertEquals(1, IntListSource.created(tag)); + + assertEquals(Arrays.asList(0, 1, 2), readBatch(1, 2)); + assertEquals("the idle reader is reused", 1, IntListSource.created(tag)); + + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals("a batch from epoch 0 is a fresh start", 2, IntListSource.created(tag)); + assertEquals(noPositions(), finalized()); + } + + /** A failed mark write fails the batch instead of advancing the reader silently. */ + @Test + public void testMarkWriteFailureFailsBatch() throws Exception { + String file = temp.newFile("not-a-directory").getAbsolutePath(); + BeamPartitionReader reader = new BeamPartitionReader<>(partition(file, 0, 1)); + List values = new ArrayList<>(); + assertThrows(IOException.class, () -> drainInto(reader, values)); + assertEquals(Arrays.asList(0, 1, 2), values); + } + + /** A retry after a failed mark write recreates the reader from the durable mark at S. */ + @Test + public void testRetryAfterFailedMarkWriteRecreatesReader() throws Exception { + String file = temp.newFile("not-a-directory").getAbsolutePath(); + assertEquals( + Arrays.asList(0, 1, 2), drainUntilFailure(partition(file, 0, 1), IOException.class)); + assertEquals( + Arrays.asList(0, 1, 2), drainUntilFailure(partition(file, 0, 1), IOException.class)); + assertEquals(noPositions(), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + /** A retry after a failed mark encode above epoch 0 restores from the durable mark at S. */ + @Test + public void testRetryAfterFailedMarkEncodeRestoresFromDurableMark() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + IntListSource.failEncoding(tag, true); + assertEquals( + Arrays.asList(3, 4, 5), drainUntilFailure(partition(1, 2), IllegalStateException.class)); + IntListSource.failEncoding(tag, false); + + assertEquals(Arrays.asList(3, 4, 5), readBatch(1, 2)); + assertEquals( + "the epoch 1 mark was finalized by the first open only", positions(3), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + // --------------------------------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------------------------------- + + private BeamInputPartition partition(long start, long end) { + return partition(location, start, end); + } + + private BeamInputPartition partition(String checkpointLocation, long start, long end) { + return new BeamInputPartition<>( + source, + CODER, + options, + hadoopConf, + checkpointLocation, + 0, + start, + end, + MAX_RECORDS, + MAX_BATCH_MILLIS, + READER_IDLE_MILLIS, + new String[0]); + } + + private BeamPartitionReader open(long start, long end) throws IOException { + return new BeamPartitionReader<>(partition(start, end)); + } + + private List readBatch(long start, long end) throws IOException { + return drain(open(start, end)); + } + + private static List drain(BeamPartitionReader reader) throws IOException { + List values = new ArrayList<>(); + drainInto(reader, values); + return values; + } + + private static void drainInto(BeamPartitionReader reader, List values) + throws IOException { + while (reader.next()) { + InternalRow row = reader.get(); + values.add(CoderUtils.decodeFromByteArray(CODER, row.getBinary(0)).getValue()); + } + reader.close(); + } + + /** Opens and drains a batch expected to fail, returns what it delivered before failing. */ + private static List drainUntilFailure( + BeamInputPartition partition, Class failure) + throws IOException { + BeamPartitionReader reader = new BeamPartitionReader<>(partition); + List values = new ArrayList<>(); + assertThrows(failure, () -> drainInto(reader, values)); + return values; + } + + private List finalized() { + return IntListSource.finalized(tag); + } + + private static List positions(Integer... positions) { + return Arrays.asList(positions); + } + + private static List noPositions() { + return Collections.emptyList(); + } + + // --------------------------------------------------------------------------------------------- + // a list source whose marks count their finalizations + // --------------------------------------------------------------------------------------------- + + /** + * Single split source over the integers {@code 0..count-1}. Marks are not Java serializable and + * record the position they finalize under the source's tag, readers are counted per tag. + */ + static class IntListSource extends UnboundedSource { + private static final long serialVersionUID = 1L; + + static final Coder MARK_CODER = new MarkCoder(); + + private static final ConcurrentMap> FINALIZED = new ConcurrentHashMap<>(); + private static final ConcurrentMap CREATED = new ConcurrentHashMap<>(); + private static final Set FAILING_ENCODE = ConcurrentHashMap.newKeySet(); + + private final String tag; + private final int count; + + IntListSource(String tag, int count) { + this.tag = tag; + this.count = count; + } + + static List finalized(String tag) { + List positions = FINALIZED.get(tag); + if (positions == null) { + return Collections.emptyList(); + } + synchronized (positions) { + return new ArrayList<>(positions); + } + } + + static int created(String tag) { + AtomicInteger created = CREATED.get(tag); + return created == null ? 0 : created.get(); + } + + static void failEncoding(String tag, boolean fail) { + if (fail) { + FAILING_ENCODE.add(tag); + } else { + FAILING_ENCODE.remove(tag); + } + } + + static void forget(String tag) { + FINALIZED.remove(tag); + CREATED.remove(tag); + FAILING_ENCODE.remove(tag); + } + + @Override + public List split(int desiredNumSplits, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark mark) { + CREATED.computeIfAbsent(tag, t -> new AtomicInteger()).incrementAndGet(); + return new IntListReader(this, mark == null ? 0 : mark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return MARK_CODER; + } + + @Override + public Coder getOutputCoder() { + return VarIntCoder.of(); + } + + /** Position of the next element, deliberately not {@link java.io.Serializable}. */ + static final class Mark implements UnboundedSource.CheckpointMark { + private final String tag; + private final int next; + + Mark(String tag, int next) { + this.tag = tag; + this.next = next; + } + + @Override + public void finalizeCheckpoint() { + FINALIZED + .computeIfAbsent(tag, t -> Collections.synchronizedList(new ArrayList<>())) + .add(next); + } + } + + private static final class MarkCoder extends CustomCoder { + private static final long serialVersionUID = 1L; + + @Override + public void encode(Mark mark, OutputStream out) throws IOException { + if (FAILING_ENCODE.contains(mark.tag)) { + throw new CoderException("injected mark encode failure for " + mark.tag); + } + StringUtf8Coder.of().encode(mark.tag, out); + VarIntCoder.of().encode(mark.next, out); + } + + @Override + public Mark decode(InputStream in) throws IOException { + return new Mark(StringUtf8Coder.of().decode(in), VarIntCoder.of().decode(in)); + } + } + + private static final class IntListReader extends UnboundedReader { + private final IntListSource source; + private int next; + private int current = -1; + + IntListReader(IntListSource source, int next) { + this.source = source; + this.next = next; + } + + @Override + public boolean start() { + return advance(); + } + + @Override + public boolean advance() { + if (next < source.count) { + current = next++; + return true; + } + return false; + } + + @Override + public Integer getCurrent() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return current; + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant(1_700_000_000_000L + current * 1_000L); + } + + @Override + public Instant getWatermark() { + return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE : getCurrentTimestamp(); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new Mark(source.tag, next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() {} + } + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java new file mode 100644 index 000000000000..e65918f219cf --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java @@ -0,0 +1,209 @@ +/* + * 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.io.streaming; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.fail; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.TreeSet; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.hadoop.conf.Configuration; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link BeamSourceCheckpoint} on the local file system, no Spark session. */ +@RunWith(JUnit4.class) +public class BeamSourceCheckpointTest { + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + private File root; + private BeamSourceCheckpoint checkpoint; + + @Before + public void setUp() { + root = new File(temp.getRoot(), "sources/0"); + checkpoint = fresh(); + } + + /** Splits round trip once, a second pin fails, a fresh location has none. */ + @Test + public void testSplitsArePinnedOnce() throws Exception { + assertNull("fresh location must have no splits", checkpoint.readSplits()); + + List> splits = + CountingSource.unbounded().split(2, PipelineOptionsFactory.create()); + assertEquals(2, splits.size()); + checkpoint.writeSplits(splits); + + List> restored = fresh().readSplits(); + assertNotNull(restored); + assertEquals(splits.size(), restored.size()); + for (int i = 0; i < splits.size(); i++) { + assertEquals(splits.get(i).getClass(), restored.get(i).getClass()); + } + + assertThrows(IOException.class, () -> checkpoint.writeSplits(splits)); + assertThrows(IOException.class, () -> fresh().writeSplits(splits)); + assertEquals(2, fresh().readSplits().size()); + assertCleanTree(); + } + + /** Marks round trip byte for byte, the latest write at an epoch wins, absent marks are null. */ + @Test + public void testMarkRoundTrip() throws Exception { + byte[] first = {0, 1, -1, 127, -128, 42}; + byte[] second = {7}; + + assertNull(checkpoint.readMark(0, 1)); + checkpoint.writeMark(0, 1, first); + assertArrayEquals(first, checkpoint.readMark(0, 1)); + assertArrayEquals(first, fresh().readMark(0, 1)); + + checkpoint.writeMark(0, 1, second); + assertArrayEquals(second, checkpoint.readMark(0, 1)); + assertArrayEquals(second, fresh().readMark(0, 1)); + + checkpoint.writeMark(0, 2, new byte[0]); + assertArrayEquals(new byte[0], checkpoint.readMark(0, 2)); + + assertNull("missing epoch", checkpoint.readMark(0, 3)); + assertNull("missing split", checkpoint.readMark(1, 1)); + assertNull("missing split, fresh instance", fresh().readMark(1, 1)); + assertCleanTree(); + } + + /** Purging keeps epochs at or above the floor, is idempotent and ignores absent splits. */ + @Test + public void testPurgeMarksBelow() throws Exception { + for (long epoch = 1; epoch <= 5; epoch++) { + checkpoint.writeMark(0, epoch, bytes(epoch)); + } + for (long epoch = 1; epoch <= 3; epoch++) { + checkpoint.writeMark(1, epoch, bytes(epoch)); + } + + checkpoint.purgeMarksBelow(0, 3); + assertEquals(epochs(3, 4, 5), markEpochs(0)); + assertEquals("other split untouched", epochs(1, 2, 3), markEpochs(1)); + + checkpoint.purgeMarksBelow(0, 3); + assertEquals("idempotent", epochs(3, 4, 5), markEpochs(0)); + + checkpoint.purgeMarksBelow(0, 1); + assertEquals("lower floor deletes nothing", epochs(3, 4, 5), markEpochs(0)); + + checkpoint.purgeMarksBelow(0, 5); + assertEquals(epochs(5), markEpochs(0)); + assertArrayEquals(bytes(5), checkpoint.readMark(0, 5)); + + checkpoint.purgeMarksBelow(0, 6); + assertEquals(epochs(), markEpochs(0)); + + checkpoint.purgeMarksBelow(7, 3); + checkpoint.purgeMarksBelow(7, 3); + assertEquals(epochs(), markEpochs(7)); + assertCleanTree(); + } + + /** A fresh instance without an in memory floor purges the same files as the writer. */ + @Test + public void testPurgeWithoutInMemoryFloor() throws Exception { + for (long epoch = 1; epoch <= 5; epoch++) { + checkpoint.writeMark(0, epoch, bytes(epoch)); + } + checkpoint.purgeMarksBelow(0, 2); + assertEquals(epochs(2, 3, 4, 5), markEpochs(0)); + + BeamSourceCheckpoint other = fresh(); + other.purgeMarksBelow(0, 4); + assertEquals(epochs(4, 5), markEpochs(0)); + other.purgeMarksBelow(0, 4); + assertEquals(epochs(4, 5), markEpochs(0)); + + // The writer's floor is 2, epochs 2 and 3 are already gone. + checkpoint.purgeMarksBelow(0, 5); + assertEquals(epochs(5), markEpochs(0)); + + fresh().purgeMarksBelow(0, 5); + assertEquals(epochs(5), markEpochs(0)); + assertCleanTree(); + } + + private BeamSourceCheckpoint fresh() { + return new BeamSourceCheckpoint(root.getAbsolutePath(), new Configuration()); + } + + private static byte[] bytes(long epoch) { + return new byte[] {(byte) epoch, (byte) (epoch * 3)}; + } + + private static Set epochs(long... epochs) { + return Arrays.stream(epochs).boxed().collect(Collectors.toCollection(TreeSet::new)); + } + + /** Numeric file names under {@code marks/}, empty if the directory is absent. */ + private Set markEpochs(int split) { + File dir = new File(root, "marks/" + split); + String[] names = dir.list(); + Set epochs = new TreeSet<>(); + if (names == null) { + return epochs; + } + for (String name : names) { + if (!name.startsWith(".")) { + epochs.add(Long.parseLong(name)); + } + } + return epochs; + } + + /** No temp files and no hidden files other than Hadoop's checksum sidecars remain. */ + private void assertCleanTree() throws IOException { + List offenders = new ArrayList<>(); + try (Stream paths = Files.walk(temp.getRoot().toPath())) { + paths + .map(path -> path.getFileName().toString()) + .filter(name -> name.endsWith(".tmp") || (name.startsWith(".") && !name.endsWith(".crc"))) + .forEach(offenders::add); + } + if (!offenders.isEmpty()) { + fail("unexpected files under " + temp.getRoot() + ": " + offenders); + } + } +} diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java index 391350fd348e..fb0192dba868 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java @@ -63,6 +63,15 @@ public interface SparkStructuredStreamingPipelineOptions extends SparkCommonPipe void setMaxBatchDurationMillis(long value); + @Description( + "Idle time in milliseconds after which an executor closes a cached unbounded reader. Must " + + "exceed the longest gap between two micro-batches, a closed reader's last checkpoint " + + "mark is not finalized and the source redelivers (streaming mode only).") + @Default.Long(600_000) + long getReaderIdleTimeoutMillis(); + + void setReaderIdleTimeoutMillis(long value); + @Description( "Test-oriented: gracefully stop streaming queries after this many consecutive empty " + "micro-batches. Disabled if negative (streaming mode only).")