| /* |
| * 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.sdk.testing; |
| |
| import static com.google.common.base.Preconditions.checkNotNull; |
| import static org.hamcrest.Matchers.containsInAnyOrder; |
| import static org.hamcrest.Matchers.equalTo; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertThat; |
| import static org.junit.Assert.assertTrue; |
| |
| import com.google.common.collect.ImmutableList; |
| import java.io.IOException; |
| import java.util.ArrayList; |
| import java.util.List; |
| import java.util.NoSuchElementException; |
| import java.util.Objects; |
| import java.util.concurrent.Callable; |
| import java.util.concurrent.CountDownLatch; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.Future; |
| import javax.annotation.Nullable; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.io.BoundedSource; |
| import org.apache.beam.sdk.io.BoundedSource.BoundedReader; |
| import org.apache.beam.sdk.io.Source; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.values.KV; |
| import org.joda.time.Instant; |
| import org.junit.Assert; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * Helper functions and test harnesses for checking correctness of {@link Source} |
| * implementations. |
| * |
| * <p>Contains a few lightweight utilities (e.g. reading items from a source or a reader, |
| * such as {@link #readFromSource} and {@link #readFromUnstartedReader}), as well as |
| * heavyweight property testing and stress testing harnesses that help getting a large |
| * amount of test coverage with few code. Most notable ones are: |
| * <ul> |
| * <li>{@link #assertSourcesEqualReferenceSource} helps testing that the data read |
| * by the union of sources produced by {@link BoundedSource#split} |
| * is the same as data read by the original source. |
| * <li>If your source implements dynamic work rebalancing, use the |
| * {@code assertSplitAtFraction} family of functions - they test behavior of |
| * {@link BoundedSource.BoundedReader#splitAtFraction}, in particular, that |
| * various consistency properties are respected and the total set of data read |
| * by the source is preserved when splits happen. |
| * Use {@link #assertSplitAtFractionBehavior} to test individual cases |
| * of {@code splitAtFraction} and use {@link #assertSplitAtFractionExhaustive} |
| * as a heavy-weight stress test including concurrency. We strongly recommend to |
| * use both. |
| * </ul> |
| * For example usages, see the unit tests of classes such as {@code AvroSource} or |
| * {@code TextSource}. |
| * |
| * <p>Like {@link PAssert}, requires JUnit and Hamcrest to be present in the classpath. |
| */ |
| public class SourceTestUtils { |
| private static final Logger LOG = LoggerFactory.getLogger(SourceTestUtils.class); |
| |
| // A wrapper around a value of type T that compares according to the structural |
| // value provided by a Coder<T>, but prints both the original and structural value, |
| // to help get good error messages from JUnit equality assertion failures and such. |
| private static class ReadableStructuralValue<T> { |
| private T originalValue; |
| private Object structuralValue; |
| |
| public ReadableStructuralValue(T originalValue, Object structuralValue) { |
| this.originalValue = originalValue; |
| this.structuralValue = structuralValue; |
| } |
| |
| @Override |
| public int hashCode() { |
| return Objects.hashCode(structuralValue); |
| } |
| |
| @Override |
| public boolean equals(Object obj) { |
| if (obj == null || !(obj instanceof ReadableStructuralValue)) { |
| return false; |
| } |
| return Objects.equals(structuralValue, ((ReadableStructuralValue) obj).structuralValue); |
| } |
| |
| @Override |
| public String toString() { |
| return String.format("[%s (structural %s)]", originalValue, structuralValue); |
| } |
| } |
| |
| /** |
| * Testing utilities below depend on standard assertions and matchers to compare elements read by |
| * sources. In general the elements may not implement {@code equals}/{@code hashCode} properly, |
| * however every source has a {@link Coder} and every {@code Coder} can |
| * produce a {@link Coder#structuralValue} whose {@code equals}/{@code hashCode} is |
| * consistent with equality of encoded format. |
| * So we use this {@link Coder#structuralValue} to compare elements read by sources. |
| */ |
| public static <T> List<ReadableStructuralValue<T>> createStructuralValues( |
| Coder<T> coder, List<T> list) |
| throws Exception { |
| List<ReadableStructuralValue<T>> result = new ArrayList<>(); |
| for (T elem : list) { |
| result.add(new ReadableStructuralValue<>(elem, coder.structuralValue(elem))); |
| } |
| return result; |
| } |
| |
| /** |
| * Reads all elements from the given {@link BoundedSource}. |
| */ |
| public static <T> List<T> readFromSource(BoundedSource<T> source, PipelineOptions options) |
| throws IOException { |
| try (BoundedSource.BoundedReader<T> reader = source.createReader(options)) { |
| return readFromUnstartedReader(reader); |
| } |
| } |
| |
| /** |
| * Reads all elements from the given unstarted {@link Source.Reader}. |
| */ |
| public static <T> List<T> readFromUnstartedReader(Source.Reader<T> reader) throws IOException { |
| return readRemainingFromReader(reader, false); |
| } |
| |
| /** |
| * Reads all elements from the given started {@link Source.Reader}. |
| */ |
| public static <T> List<T> readFromStartedReader(Source.Reader<T> reader) throws IOException { |
| return readRemainingFromReader(reader, true); |
| } |
| |
| /** |
| * Read elements from a {@link Source.Reader} until n elements are read. |
| */ |
| public static <T> List<T> readNItemsFromUnstartedReader(Source.Reader<T> reader, int n) |
| throws IOException { |
| return readNItemsFromReader(reader, n, false); |
| } |
| |
| /** |
| * Read elements from a {@link Source.Reader} that has already had {@link Source.Reader#start} |
| * called on it, until n elements are read. |
| */ |
| public static <T> List<T> readNItemsFromStartedReader(Source.Reader<T> reader, int n) |
| throws IOException { |
| return readNItemsFromReader(reader, n, true); |
| } |
| |
| /** |
| * Read elements from a {@link Source.Reader} until n elements are read. |
| * |
| * <p>There must be at least n elements remaining in the reader, except for |
| * the case when n is {@code Integer.MAX_VALUE}, which means "read all |
| * remaining elements". |
| */ |
| private static <T> List<T> readNItemsFromReader(Source.Reader<T> reader, int n, boolean started) |
| throws IOException { |
| List<T> res = new ArrayList<>(); |
| for (int i = 0; i < n; i++) { |
| boolean shouldStart = (i == 0 && !started); |
| boolean more = shouldStart ? reader.start() : reader.advance(); |
| if (n != Integer.MAX_VALUE) { |
| assertTrue(more); |
| } |
| if (!more) { |
| break; |
| } |
| res.add(reader.getCurrent()); |
| } |
| return res; |
| } |
| |
| /** |
| * Read all remaining elements from a {@link Source.Reader}. |
| */ |
| public static <T> List<T> readRemainingFromReader(Source.Reader<T> reader, boolean started) |
| throws IOException { |
| return readNItemsFromReader(reader, Integer.MAX_VALUE, started); |
| } |
| |
| /** |
| * Given a reference {@code Source} and a list of {@code Source}s, assert that the union of |
| * the records read from the list of sources is equal to the records read from the reference |
| * source. |
| */ |
| public static <T> void assertSourcesEqualReferenceSource( |
| BoundedSource<T> referenceSource, |
| List<? extends BoundedSource<T>> sources, |
| PipelineOptions options) |
| throws Exception { |
| Coder<T> coder = referenceSource.getDefaultOutputCoder(); |
| List<T> referenceRecords = readFromSource(referenceSource, options); |
| List<T> bundleRecords = new ArrayList<>(); |
| for (BoundedSource<T> source : sources) { |
| assertThat( |
| "Coder type for source " |
| + source |
| + " is not compatible with Coder type for referenceSource " |
| + referenceSource, |
| source.getDefaultOutputCoder(), |
| equalTo(coder)); |
| List<T> elems = readFromSource(source, options); |
| bundleRecords.addAll(elems); |
| } |
| List<ReadableStructuralValue<T>> bundleValues = |
| createStructuralValues(coder, bundleRecords); |
| List<ReadableStructuralValue<T>> referenceValues = |
| createStructuralValues(coder, referenceRecords); |
| assertThat(bundleValues, containsInAnyOrder(referenceValues.toArray())); |
| } |
| |
| /** |
| * Assert that a {@code Reader} returns a {@code Source} that, when read from, produces the same |
| * records as the reader. |
| */ |
| public static <T> void assertUnstartedReaderReadsSameAsItsSource( |
| BoundedSource.BoundedReader<T> reader, PipelineOptions options) throws Exception { |
| Coder<T> coder = reader.getCurrentSource().getDefaultOutputCoder(); |
| List<T> expected = readFromUnstartedReader(reader); |
| List<T> actual = readFromSource(reader.getCurrentSource(), options); |
| List<ReadableStructuralValue<T>> expectedStructural = createStructuralValues(coder, expected); |
| List<ReadableStructuralValue<T>> actualStructural = createStructuralValues(coder, actual); |
| assertThat(actualStructural, containsInAnyOrder(expectedStructural.toArray())); |
| } |
| |
| /** |
| * Expected outcome of |
| * {@link org.apache.beam.sdk.io.BoundedSource.BoundedReader#splitAtFraction}. |
| */ |
| public enum ExpectedSplitOutcome { |
| /** |
| * The operation must succeed and the results must be consistent. |
| */ |
| MUST_SUCCEED_AND_BE_CONSISTENT, |
| /** |
| * The operation must fail (return {@code null}). |
| */ |
| MUST_FAIL, |
| /** |
| * The operation must either fail, or succeed and the results be consistent. |
| */ |
| MUST_BE_CONSISTENT_IF_SUCCEEDS |
| } |
| |
| /** |
| * Contains two values: the number of items in the primary source, and the number of items in |
| * the residual source, -1 if split failed. |
| */ |
| private static class SplitAtFractionResult { |
| public int numPrimaryItems; |
| public int numResidualItems; |
| |
| public SplitAtFractionResult(int numPrimaryItems, int numResidualItems) { |
| this.numPrimaryItems = numPrimaryItems; |
| this.numResidualItems = numResidualItems; |
| } |
| } |
| |
| /** |
| * Asserts that the {@code source}'s reader either fails to {@code splitAtFraction(fraction)} |
| * after reading {@code numItemsToReadBeforeSplit} items, or succeeds in a way that is |
| * consistent according to {@link #assertSplitAtFractionSucceedsAndConsistent}. |
| * |
| * <p>Returns SplitAtFractionResult. |
| */ |
| |
| public static <T> SplitAtFractionResult assertSplitAtFractionBehavior( |
| BoundedSource<T> source, |
| int numItemsToReadBeforeSplit, |
| double splitFraction, |
| ExpectedSplitOutcome expectedOutcome, |
| PipelineOptions options) |
| throws Exception { |
| return assertSplitAtFractionBehaviorImpl( |
| source, readFromSource(source, options), numItemsToReadBeforeSplit, splitFraction, |
| expectedOutcome, options); |
| } |
| |
| /** |
| * Compares two lists elementwise and throws a detailed assertion failure optimized for |
| * human reading in case they are unequal. |
| */ |
| private static <T> void assertListsEqualInOrder( |
| String message, String expectedLabel, List<T> expected, String actualLabel, List<T> actual) { |
| int i = 0; |
| for (; i < expected.size() && i < actual.size(); ++i) { |
| if (!Objects.equals(expected.get(i), actual.get(i))) { |
| Assert.fail(String.format( |
| "%s: %s and %s have %d items in common and then differ. " |
| + "Item in %s (%d more): %s, item in %s (%d more): %s", |
| message, expectedLabel, actualLabel, i, |
| expectedLabel, expected.size() - i - 1, expected.get(i), |
| actualLabel, actual.size() - i - 1, actual.get(i))); |
| } |
| } |
| if (i < expected.size() /* but i == actual.size() */) { |
| Assert.fail(String.format( |
| "%s: %s has %d more items after matching all %d from %s. First 5: %s", |
| message, expectedLabel, expected.size() - actual.size(), actual.size(), actualLabel, |
| expected.subList(actual.size(), Math.min(expected.size(), actual.size() + 5)))); |
| } else if (i < actual.size() /* but i == expected.size() */) { |
| Assert.fail(String.format( |
| "%s: %s has %d more items after matching all %d from %s. First 5: %s", |
| message, actualLabel, actual.size() - expected.size(), expected.size(), expectedLabel, |
| actual.subList(expected.size(), Math.min(actual.size(), expected.size() + 5)))); |
| } else { |
| // All is well. |
| } |
| } |
| |
| private static <T> SourceTestUtils.SplitAtFractionResult assertSplitAtFractionBehaviorImpl( |
| BoundedSource<T> source, List<T> expectedItems, int numItemsToReadBeforeSplit, |
| double splitFraction, ExpectedSplitOutcome expectedOutcome, PipelineOptions options) |
| throws Exception { |
| try (BoundedSource.BoundedReader<T> reader = source.createReader(options)) { |
| BoundedSource<T> originalSource = reader.getCurrentSource(); |
| List<T> currentItems = readNItemsFromUnstartedReader(reader, numItemsToReadBeforeSplit); |
| BoundedSource<T> residual = reader.splitAtFraction(splitFraction); |
| if (residual != null) { |
| assertFalse( |
| String.format( |
| "Primary source didn't change after a successful split of %s at %f " |
| + "after reading %d items. " |
| + "Was the source object mutated instead of creating a new one? " |
| + "Source objects MUST be immutable.", |
| source, splitFraction, numItemsToReadBeforeSplit), |
| reader.getCurrentSource() == originalSource); |
| assertFalse( |
| String.format( |
| "Residual source equal to original source after a successful split of %s at %f " |
| + "after reading %d items. " |
| + "Was the source object mutated instead of creating a new one? " |
| + "Source objects MUST be immutable.", |
| source, splitFraction, numItemsToReadBeforeSplit), |
| reader.getCurrentSource() == residual); |
| } |
| // Failure cases are: must succeed but fails; must fail but succeeds. |
| switch (expectedOutcome) { |
| case MUST_SUCCEED_AND_BE_CONSISTENT: |
| assertNotNull( |
| "Failed to split reader of source: " |
| + source |
| + " at " |
| + splitFraction |
| + " after reading " |
| + numItemsToReadBeforeSplit |
| + " items", |
| residual); |
| break; |
| case MUST_FAIL: |
| assertEquals(null, residual); |
| break; |
| case MUST_BE_CONSISTENT_IF_SUCCEEDS: |
| // Nothing. |
| break; |
| } |
| currentItems.addAll(readRemainingFromReader(reader, numItemsToReadBeforeSplit > 0)); |
| BoundedSource<T> primary = reader.getCurrentSource(); |
| return verifySingleSplitAtFractionResult( |
| source, expectedItems, currentItems, primary, residual, |
| numItemsToReadBeforeSplit, splitFraction, options); |
| } |
| } |
| |
| private static <T> SourceTestUtils.SplitAtFractionResult verifySingleSplitAtFractionResult( |
| BoundedSource<T> source, List<T> expectedItems, List<T> currentItems, |
| BoundedSource<T> primary, BoundedSource<T> residual, |
| int numItemsToReadBeforeSplit, double splitFraction, PipelineOptions options) |
| throws Exception { |
| List<T> primaryItems = readFromSource(primary, options); |
| if (residual != null) { |
| List<T> residualItems = readFromSource(residual, options); |
| List<T> totalItems = new ArrayList<>(); |
| totalItems.addAll(primaryItems); |
| totalItems.addAll(residualItems); |
| String errorMsgForPrimarySourceComp = |
| String.format( |
| "Continued reading after split yielded different items than primary source: " |
| + "split at %s after reading %s items, original source: %s, primary source: %s", |
| splitFraction, |
| numItemsToReadBeforeSplit, |
| source, |
| primary); |
| String errorMsgForTotalSourceComp = |
| String.format( |
| "Items in primary and residual sources after split do not add up to items " |
| + "in the original source. Split at %s after reading %s items; " |
| + "original source: %s, primary: %s, residual: %s", |
| splitFraction, |
| numItemsToReadBeforeSplit, |
| source, |
| primary, |
| residual); |
| Coder<T> coder = primary.getDefaultOutputCoder(); |
| List<ReadableStructuralValue<T>> primaryValues = |
| createStructuralValues(coder, primaryItems); |
| List<ReadableStructuralValue<T>> currentValues = |
| createStructuralValues(coder, currentItems); |
| List<ReadableStructuralValue<T>> expectedValues = |
| createStructuralValues(coder, expectedItems); |
| List<ReadableStructuralValue<T>> totalValues = |
| createStructuralValues(coder, totalItems); |
| assertListsEqualInOrder( |
| errorMsgForPrimarySourceComp, "current", currentValues, "primary", primaryValues); |
| assertListsEqualInOrder( |
| errorMsgForTotalSourceComp, "total", expectedValues, "primary+residual", totalValues); |
| return new SplitAtFractionResult(primaryItems.size(), residualItems.size()); |
| } |
| return new SplitAtFractionResult(primaryItems.size(), -1); |
| } |
| |
| /** |
| * Verifies some consistency properties of |
| * {@link BoundedSource.BoundedReader#splitAtFraction} on the given source. Equivalent to |
| * the following pseudocode: |
| * <pre> |
| * Reader reader = source.createReader(); |
| * read N items from reader; |
| * Source residual = reader.splitAtFraction(splitFraction); |
| * Source primary = reader.getCurrentSource(); |
| * assert: items in primary == items we read so far |
| * + items we'll get by continuing to read from reader; |
| * assert: items in original source == items in primary + items in residual |
| * </pre> |
| */ |
| public static <T> void assertSplitAtFractionSucceedsAndConsistent( |
| BoundedSource<T> source, |
| int numItemsToReadBeforeSplit, |
| double splitFraction, |
| PipelineOptions options) |
| throws Exception { |
| assertSplitAtFractionBehavior( |
| source, |
| numItemsToReadBeforeSplit, |
| splitFraction, |
| ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT, |
| options); |
| } |
| |
| /** |
| * Asserts that the {@code source}'s reader fails to {@code splitAtFraction(fraction)} |
| * after reading {@code numItemsToReadBeforeSplit} items. |
| */ |
| public static <T> void assertSplitAtFractionFails( |
| BoundedSource<T> source, |
| int numItemsToReadBeforeSplit, |
| double splitFraction, |
| PipelineOptions options) |
| throws Exception { |
| assertSplitAtFractionBehavior( |
| source, numItemsToReadBeforeSplit, splitFraction, ExpectedSplitOutcome.MUST_FAIL, options); |
| } |
| |
| private static class SplitFractionStatistics { |
| List<Double> successfulFractions = new ArrayList<>(); |
| List<Double> nonTrivialFractions = new ArrayList<>(); |
| } |
| |
| /** |
| * Asserts that given a start position, |
| * {@link BoundedSource.BoundedReader#splitAtFraction} at every interesting fraction (halfway |
| * between two fractions that differ by at least one item) can be called successfully and the |
| * results are consistent if a split succeeds. |
| */ |
| private static <T> void assertSplitAtFractionBinary( |
| BoundedSource<T> source, |
| List<T> expectedItems, |
| int numItemsToBeReadBeforeSplit, |
| double leftFraction, |
| SplitAtFractionResult leftResult, |
| double rightFraction, |
| SplitAtFractionResult rightResult, |
| PipelineOptions options, |
| SplitFractionStatistics stats) |
| throws Exception { |
| if (rightFraction - leftFraction < 0.001) { |
| // Do not recurse too deeply. Otherwise we will end up in infinite |
| // recursion, e.g., while trying to find the exact minimal fraction s.t. |
| // split succeeds. A precision of 0.001 when looking for such a fraction |
| // ought to be enough for everybody. |
| return; |
| } |
| double middleFraction = (rightFraction + leftFraction) / 2; |
| if (leftResult == null) { |
| leftResult = assertSplitAtFractionBehaviorImpl( |
| source, expectedItems, numItemsToBeReadBeforeSplit, leftFraction, |
| ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options); |
| } |
| if (rightResult == null) { |
| rightResult = assertSplitAtFractionBehaviorImpl( |
| source, expectedItems, numItemsToBeReadBeforeSplit, rightFraction, |
| ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options); |
| } |
| SplitAtFractionResult middleResult = assertSplitAtFractionBehaviorImpl( |
| source, expectedItems, numItemsToBeReadBeforeSplit, middleFraction, |
| ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options); |
| if (middleResult.numResidualItems != -1) { |
| stats.successfulFractions.add(middleFraction); |
| } |
| if (middleResult.numResidualItems > 0) { |
| stats.nonTrivialFractions.add(middleFraction); |
| } |
| // Two split fractions are equivalent if they yield the same number of |
| // items in primary vs. residual source. Left and right are already not |
| // equivalent. Recurse into [left, middle) and [right, middle) respectively |
| // if middle is not equivalent to left or right. |
| if (leftResult.numPrimaryItems != middleResult.numPrimaryItems) { |
| assertSplitAtFractionBinary( |
| source, expectedItems, numItemsToBeReadBeforeSplit, |
| leftFraction, leftResult, middleFraction, middleResult, options, stats); |
| } |
| if (rightResult.numPrimaryItems != middleResult.numPrimaryItems) { |
| assertSplitAtFractionBinary( |
| source, expectedItems, numItemsToBeReadBeforeSplit, |
| middleFraction, middleResult, rightFraction, rightResult, options, stats); |
| } |
| } |
| |
| private static final int MAX_CONCURRENT_SPLITTING_TRIALS_PER_ITEM = 100; |
| private static final int MAX_CONCURRENT_SPLITTING_TRIALS_TOTAL = 1000; |
| |
| /** |
| * Asserts that for each possible start position, |
| * {@link BoundedSource.BoundedReader#splitAtFraction} at every interesting fraction (halfway |
| * between two fractions that differ by at least one item) can be called successfully and the |
| * results are consistent if a split succeeds. Verifies multithreaded splitting as well. |
| */ |
| public static <T> void assertSplitAtFractionExhaustive( |
| BoundedSource<T> source, PipelineOptions options) throws Exception { |
| List<T> expectedItems = readFromSource(source, options); |
| assertFalse("Empty source", expectedItems.isEmpty()); |
| assertFalse("Source reads a single item", expectedItems.size() == 1); |
| List<List<Double>> allNonTrivialFractions = new ArrayList<>(); |
| { |
| boolean anySuccessfulFractions = false; |
| boolean anyNonTrivialFractions = false; |
| for (int i = 0; i < expectedItems.size(); i++) { |
| SplitFractionStatistics stats = new SplitFractionStatistics(); |
| assertSplitAtFractionBinary(source, expectedItems, i, |
| 0.0, null, 1.0, null, options, stats); |
| if (!stats.successfulFractions.isEmpty()) { |
| anySuccessfulFractions = true; |
| } |
| if (!stats.nonTrivialFractions.isEmpty()) { |
| anyNonTrivialFractions = true; |
| } |
| allNonTrivialFractions.add(stats.nonTrivialFractions); |
| } |
| assertTrue( |
| "splitAtFraction test completed vacuously: no successful split fractions found", |
| anySuccessfulFractions); |
| assertTrue( |
| "splitAtFraction test completed vacuously: no non-trivial split fractions found", |
| anyNonTrivialFractions); |
| } |
| { |
| // Perform a stress test of "racy" concurrent splitting: |
| // for every position (number of items read), try to split at the minimum nontrivial |
| // split fraction for that position concurrently with reading the record at that position. |
| // To ensure that the test is non-vacuous, make sure that the splitting succeeds |
| // at least once and fails at least once. |
| ExecutorService executor = Executors.newFixedThreadPool(2); |
| int numTotalTrials = 0; |
| for (int i = 0; i < expectedItems.size(); i++) { |
| double minNonTrivialFraction = 2.0; // Greater than any possible fraction. |
| for (double fraction : allNonTrivialFractions.get(i)) { |
| minNonTrivialFraction = Math.min(minNonTrivialFraction, fraction); |
| } |
| if (minNonTrivialFraction == 2.0) { |
| // This will not happen all the time because otherwise the test above would |
| // detect vacuousness. |
| continue; |
| } |
| int numTrials = 0; |
| boolean haveSuccess = false, haveFailure = false; |
| while (true) { |
| ++numTrials; |
| if (numTrials > MAX_CONCURRENT_SPLITTING_TRIALS_PER_ITEM) { |
| LOG.warn( |
| "After {} concurrent splitting trials at item #{}, observed only {}, " |
| + "giving up on this item", |
| numTrials, i, haveSuccess ? "success" : "failure"); |
| break; |
| } |
| if (assertSplitAtFractionConcurrent( |
| executor, source, expectedItems, i, minNonTrivialFraction, options)) { |
| haveSuccess = true; |
| } else { |
| haveFailure = true; |
| } |
| if (haveSuccess && haveFailure) { |
| LOG.info( |
| "{} trials to observe both success and failure of concurrent splitting at item #{}", |
| numTrials, i); |
| break; |
| } |
| } |
| numTotalTrials += numTrials; |
| if (numTotalTrials > MAX_CONCURRENT_SPLITTING_TRIALS_TOTAL) { |
| LOG.warn( |
| "After {} total concurrent splitting trials, considered only {} items, giving up.", |
| numTotalTrials, i); |
| break; |
| } |
| } |
| LOG.info( |
| "{} total concurrent splitting trials for {} items", |
| numTotalTrials, expectedItems.size()); |
| } |
| } |
| |
| private static <T> boolean assertSplitAtFractionConcurrent( |
| ExecutorService executor, BoundedSource<T> source, List<T> expectedItems, |
| final int numItemsToReadBeforeSplitting, final double fraction, PipelineOptions options) |
| throws Exception { |
| @SuppressWarnings("resource") // Closed in readerThread |
| final BoundedSource.BoundedReader<T> reader = source.createReader(options); |
| final CountDownLatch unblockSplitter = new CountDownLatch(1); |
| Future<List<T>> readerThread = |
| executor.submit( |
| new Callable<List<T>>() { |
| @Override |
| public List<T> call() throws Exception { |
| try { |
| List<T> items = |
| readNItemsFromUnstartedReader(reader, numItemsToReadBeforeSplitting); |
| unblockSplitter.countDown(); |
| items.addAll(readRemainingFromReader(reader, numItemsToReadBeforeSplitting > 0)); |
| return items; |
| } finally { |
| reader.close(); |
| } |
| } |
| }); |
| Future<KV<BoundedSource<T>, BoundedSource<T>>> splitterThread = executor.submit( |
| new Callable<KV<BoundedSource<T>, BoundedSource<T>>>() { |
| @Override |
| public KV<BoundedSource<T>, BoundedSource<T>> call() throws Exception { |
| unblockSplitter.await(); |
| BoundedSource<T> residual = reader.splitAtFraction(fraction); |
| if (residual == null) { |
| return null; |
| } |
| return KV.of(reader.getCurrentSource(), residual); |
| } |
| }); |
| List<T> currentItems = readerThread.get(); |
| KV<BoundedSource<T>, BoundedSource<T>> splitSources = splitterThread.get(); |
| if (splitSources == null) { |
| return false; |
| } |
| SplitAtFractionResult res = verifySingleSplitAtFractionResult( |
| source, expectedItems, currentItems, splitSources.getKey(), splitSources.getValue(), |
| numItemsToReadBeforeSplitting, fraction, options); |
| return (res.numResidualItems > 0); |
| } |
| |
| /** |
| * Returns an equivalent unsplittable {@code BoundedSource<T>}. |
| * |
| * <p>It forwards most methods to the given {@code boundedSource}, except: |
| * <ol> |
| * <li> {@link BoundedSource#split} rejects initial splitting |
| * by returning itself in a list. |
| * <li> {@link BoundedReader#splitAtFraction} rejects dynamic splitting by returning null. |
| * </ol> |
| */ |
| public static <T> BoundedSource<T> toUnsplittableSource(BoundedSource<T> boundedSource) { |
| return new UnsplittableSource<>(boundedSource); |
| } |
| |
| private static class UnsplittableSource<T> extends BoundedSource<T> { |
| |
| private final BoundedSource<T> boundedSource; |
| |
| private UnsplittableSource(BoundedSource<T> boundedSource) { |
| this.boundedSource = checkNotNull(boundedSource, "boundedSource"); |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| this.boundedSource.populateDisplayData(builder); |
| } |
| |
| @Override |
| public List<? extends BoundedSource<T>> split( |
| long desiredBundleSizeBytes, PipelineOptions options) throws Exception { |
| return ImmutableList.of(this); |
| } |
| |
| @Override |
| public long getEstimatedSizeBytes(PipelineOptions options) throws Exception { |
| return boundedSource.getEstimatedSizeBytes(options); |
| } |
| |
| @Override |
| public BoundedReader<T> createReader(PipelineOptions options) throws IOException { |
| return new UnsplittableReader<>(boundedSource, boundedSource.createReader(options)); |
| } |
| |
| @Override |
| public void validate() { |
| boundedSource.validate(); |
| } |
| |
| @Override |
| public Coder<T> getDefaultOutputCoder() { |
| return boundedSource.getDefaultOutputCoder(); |
| } |
| |
| private static class UnsplittableReader<T> extends BoundedReader<T> { |
| |
| private final BoundedSource<T> boundedSource; |
| private final BoundedReader<T> boundedReader; |
| |
| private UnsplittableReader(BoundedSource<T> boundedSource, BoundedReader<T> boundedReader) { |
| this.boundedSource = checkNotNull(boundedSource, "boundedSource"); |
| this.boundedReader = checkNotNull(boundedReader, "boundedReader"); |
| } |
| |
| @Override |
| public BoundedSource<T> getCurrentSource() { |
| return boundedSource; |
| } |
| |
| @Override |
| public boolean start() throws IOException { |
| return boundedReader.start(); |
| } |
| |
| @Override |
| public boolean advance() throws IOException { |
| return boundedReader.advance(); |
| } |
| |
| @Override |
| public T getCurrent() throws NoSuchElementException { |
| return boundedReader.getCurrent(); |
| } |
| |
| @Override |
| public void close() throws IOException { |
| boundedReader.close(); |
| } |
| |
| @Override |
| @Nullable |
| public BoundedSource<T> splitAtFraction(double fraction) { |
| return null; |
| } |
| |
| @Override |
| @Nullable |
| public Double getFractionConsumed() { |
| return boundedReader.getFractionConsumed(); |
| } |
| |
| @Override |
| public long getSplitPointsConsumed() { |
| return boundedReader.getSplitPointsConsumed(); |
| } |
| |
| @Override |
| public long getSplitPointsRemaining() { |
| return boundedReader.getSplitPointsRemaining(); |
| } |
| |
| @Override |
| public Instant getCurrentTimestamp() throws NoSuchElementException { |
| return boundedReader.getCurrentTimestamp(); |
| } |
| } |
| } |
| } |
