diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json index 97b7503f5445..f7bf5e8a08a0 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 7, + "modification": 14, } diff --git a/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow.yml b/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow.yml index 45231941ce02..76b842f1e11b 100644 --- a/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow.yml +++ b/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow.yml @@ -74,7 +74,7 @@ jobs: uses: ./.github/actions/gradle-command-self-hosted-action with: gradle-command: :runners:google-cloud-dataflow-java:validatesRunner - max-workers: 12 + max-workers: 24 - name: Archive JUnit Test Results uses: actions/upload-artifact@v7 if: ${{ !success() }} diff --git a/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_JavaVersions.yml b/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_JavaVersions.yml index 92639eeb10c2..995764cef7d2 100644 --- a/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_JavaVersions.yml +++ b/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_JavaVersions.yml @@ -81,7 +81,7 @@ jobs: arguments: | -PtestJavaVersion=${{ matrix.java_version }} \ -Pjava${{ matrix.java_version }}Home=$JAVA_HOME_${{ matrix.java_version }}_X64 \ - max-workers: 12 + max-workers: 24 - name: Archive JUnit Test Results uses: actions/upload-artifact@v7 if: ${{ !success() }} diff --git a/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_V2.yml b/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_V2.yml index 443f0c920d23..55922dfb3ab6 100644 --- a/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_V2.yml +++ b/.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_V2.yml @@ -74,7 +74,7 @@ jobs: uses: ./.github/actions/gradle-command-self-hosted-action with: gradle-command: :runners:google-cloud-dataflow-java:validatesRunnerV2 - max-workers: 12 + max-workers: 24 - name: Archive JUnit Test Results uses: actions/upload-artifact@v7 if: ${{ !success() }} diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle index c11a04a71979..d0aa2dfc0b10 100644 --- a/runners/google-cloud-dataflow-java/build.gradle +++ b/runners/google-cloud-dataflow-java/build.gradle @@ -132,6 +132,11 @@ dependencies { implementation library.java.joda_time implementation library.java.slf4j_api implementation library.java.vendored_grpc_1_69_0 + // TestDataflowRunner inspects TestPipeline, a JUnit TestRule, which needs JUnit at compile + // time (as in sdks/java/core and runners/direct-java). Users of TestDataflowRunner already have + // JUnit on their classpath. + provided library.java.junit + permitUnusedDeclared library.java.junit testImplementation library.java.guava_testlib testImplementation library.java.junit testImplementation project(path: ":sdks:java:io:google-cloud-platform", configuration: "testRuntimeMigration") @@ -241,13 +246,38 @@ def commonRunnerV2ExcludeCategories = [ 'org.apache.beam.sdk.testing.UsesBoundedTrieMetrics', // Dataflow QM as of now does not support returning back BoundedTrie in metric result. ] - -def validatesRunnerParallelThreads = project.findProperty('parallelThreads') ?: '6' +// ValidatesRunner parallelism. Each test worker JVM runs ValidatesRunner tests on several threads; +// where TestDataflowRunner can merge tests into one Dataflow job (-PtestBatching, batch mode only +// for now) many more threads are useful since most of them just contribute a member to the next +// merged job rather than a job of their own. Tests that cannot be merged, including every test +// of a streaming task, are still one job each, so their thread count is kept at the pre-batching +// level to bound the number of concurrent jobs per worker JVM. +def validatesRunnerTestBatching = project.findProperty('testBatching') ?: 'true' +def validatesRunnerTestBatchWindowMs = project.findProperty('testBatchWindowMs') ?: '3000' +def validatesRunnerMaxConcurrentStandaloneJobs = project.findProperty('maxConcurrentStandaloneJobs') ?: '12' + +def configureValidatesRunnerParallelism = { Test task, List pipelineOptions -> + def streaming = pipelineOptions.any { it == '--streaming' || it == '--streaming=true' } + def batching = validatesRunnerTestBatching == 'true' && !streaming + def parallelThreads = project.findProperty('parallelThreads') ?: (batching ? '30' : '6') + def batchMaxSize = project.findProperty('testBatchMaxSize') ?: parallelThreads + + task.systemProperty "beam.validatesRunner.parallelThreads", parallelThreads + task.systemProperty "beam.dataflow.testBatching", batching.toString() + task.systemProperty "beam.dataflow.testBatchMaxSize", batchMaxSize + task.systemProperty "beam.dataflow.testBatchWindowMs", validatesRunnerTestBatchWindowMs + task.systemProperty "beam.dataflow.maxConcurrentStandaloneJobs", validatesRunnerMaxConcurrentStandaloneJobs + // Merged jobs tell member tests apart by their unique TestPipeline root names. + task.systemProperty "beamTestPipelineUniqueRootNames", batching.toString() +} def isValidatesRunnerTestClass = { FileTreeElement element -> if (element.isDirectory()) { return true } + // Classes that reference the ValidatesRunner annotation but are not ValidatesRunner tests: + // BeamParallelJunit4RunnerTest's nested fixtures carry @Category(ValidatesRunner.class) to test + // the runner itself and must not be executed against a real service. if (!element.name.endsWith('.class') || element.name.startsWith('ValidateRunnerXlangTest') || element.name.startsWith('BeamParallelJunit4RunnerTest')) { @@ -278,7 +308,7 @@ def createLegacyWorkerValidatesRunnerTest = { Map args -> dependsOn ":runners:google-cloud-dataflow-java:worker:shadowJar" systemProperty "beamTestPipelineOptions", JsonOutput.toJson(pipelineOptions) - systemProperty "beam.validatesRunner.parallelThreads", validatesRunnerParallelThreads + configureValidatesRunnerParallelism(it, pipelineOptions) // Increase test parallelism up to the number of Gradle workers. By default this is equal // to the number of CPU cores, but can be increased by setting --max-workers=N. @@ -314,12 +344,14 @@ def createRunnerV2ValidatesRunnerTest = { Map args -> return tasks.create(name: name, type: Test, group: "Verification") { dependsOn buildAndPushDockerJavaContainer systemProperty "beamTestPipelineOptions", JsonOutput.toJson(pipelineOptions) - systemProperty "beam.validatesRunner.parallelThreads", validatesRunnerParallelThreads + configureValidatesRunnerParallelism(it, pipelineOptions) // Increase test parallelism up to the number of Gradle workers. By default this is equal // to the number of CPU cores, but can be increased by setting --max-workers=N. maxParallelForks Integer.MAX_VALUE - maxHeapSize = '768m' + // Same as the legacy task: in batch mode a fork hosts up to 30 concurrent test threads plus + // the translation of each merged job, not just a handful of standalone pipelines. + maxHeapSize = '2g' classpath = configurations.validatesRunner testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs) + files(project(project.path).sourceSets.test.output.classesDirs) @@ -646,8 +678,8 @@ task validatesRunnerV2 { 'org.apache.beam.sdk.transforms.ReshuffleTest.testReshuffleWithTimestampsStreaming', // TODO(https://github.com/apache/beam/issues/37371): Identify the cause of assertion failure for these tests - 'org.apache.beam.runners.dataflow.DataflowRunnerTest.testBatchGroupIntoBatchesOverrideBytes', - 'org.apache.beam.runners.dataflow.DataflowRunnerTest.testBatchGroupIntoBatchesOverrideCount', + 'org.apache.beam.runners.dataflow.DataflowRunnerValidatesRunnerTest.testBatchGroupIntoBatchesOverrideBytes', + 'org.apache.beam.runners.dataflow.DataflowRunnerValidatesRunnerTest.testBatchGroupIntoBatchesOverrideCount', // These tests use static state and don't work with remote execution. 'org.apache.beam.sdk.transforms.ParDoLifecycleTest.testTeardownCalledAfterExceptionInFinishBundle', @@ -689,7 +721,7 @@ task validatesRunnerV2Streaming { // TestStream with processing time control not supported 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testAfterProcessingTimeContinuationTriggerUsingState', - 'org.apache.beam.sdk.transforms.GroupByKeyTest.testCombiningAccumulatingProcessingTime', + 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testCombiningAccumulatingProcessingTime', // These tests use static state and don't work with remote execution. 'org.apache.beam.sdk.transforms.ParDoLifecycleTest.testTeardownCalledAfterExceptionInFinishBundle', diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowClient.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowClient.java index d643c9e40cde..933ef5aad45c 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowClient.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowClient.java @@ -94,6 +94,14 @@ public Job updateJob(@Nonnull String jobId, @Nonnull Job content) throws IOExcep /** Gets the Dataflow {@link Job} with the given {@code jobId}. */ public Job getJob(@Nonnull String jobId) throws IOException { + return getJob(jobId, null); + } + + /** + * Gets the Dataflow {@link Job} with the given {@code jobId} using the given {@code view} (one of + * the {@code JOB_VIEW_*} values), or the service default view if {@code null}. + */ + public Job getJob(@Nonnull String jobId, @Nullable String view) throws IOException { checkNotNull(jobId, "jobId"); Jobs.Get jobsGet = dataflow @@ -101,6 +109,9 @@ public Job getJob(@Nonnull String jobId) throws IOException { .locations() .jobs() .get(options.getProject(), options.getRegion(), jobId); + if (view != null) { + jobsGet.setView(view); + } return jobsGet.execute(); } diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowTestBatchCoordinator.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowTestBatchCoordinator.java new file mode 100644 index 000000000000..874f1cd202c6 --- /dev/null +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowTestBatchCoordinator.java @@ -0,0 +1,1414 @@ +/* + * 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.dataflow; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.api.services.dataflow.model.ComponentSource; +import com.google.api.services.dataflow.model.ComponentTransform; +import com.google.api.services.dataflow.model.ExecutionStageState; +import com.google.api.services.dataflow.model.ExecutionStageSummary; +import com.google.api.services.dataflow.model.Job; +import com.google.api.services.dataflow.model.JobMetrics; +import com.google.api.services.dataflow.model.StageSource; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.dataflow.TestDataflowRunner.ErrorMonitorMessagesHandler; +import org.apache.beam.runners.dataflow.TestDataflowRunner.PAssertCounts; +import org.apache.beam.runners.dataflow.util.MonitoringUtil; +import org.apache.beam.sdk.CompositePipeline; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.Pipeline.PipelineVisitor; +import org.apache.beam.sdk.PipelineResult.State; +import org.apache.beam.sdk.metrics.MetricFiltering; +import org.apache.beam.sdk.metrics.MetricKey; +import org.apache.beam.sdk.metrics.MetricNameFilter; +import org.apache.beam.sdk.metrics.MetricQueryResults; +import org.apache.beam.sdk.metrics.MetricResult; +import org.apache.beam.sdk.metrics.MetricResults; +import org.apache.beam.sdk.metrics.MetricsFilter; +import org.apache.beam.sdk.options.ApplicationNameOptions; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.runners.AppliedPTransform; +import org.apache.beam.sdk.runners.TransformHierarchy; +import org.apache.beam.sdk.testing.BeamParallelJunit4Runner; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.SerializableMatchers; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestPipelineOptions; +import org.apache.beam.sdk.util.HistogramData; +import org.apache.beam.sdk.util.common.ReflectHelpers; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollection.IsBounded; +import org.apache.beam.sdk.values.PValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.BiMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.DateTimeZone; +import org.joda.time.Duration; +import org.joda.time.format.DateTimeFormat; +import org.joda.time.format.DateTimeFormatter; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Coordinates merging concurrent batch {@code TestPipeline} executions within a JVM into a single + * Dataflow job per batch, with scoped {@link PAssert} and {@link MetricResults} verification. + * + *

Batching spans runner instances: every test builds its own {@link TestDataflowRunner}, and the + * pipelines merged into one job come from different tests, threads and classes. Runners therefore + * share the single process-wide coordinator returned by {@link #shared()}; tests may construct + * private instances instead. + * + *

Scoping: members of a merged job are told apart by their {@linkplain TestPipeline#getRootName + * unique root name}, which {@link TestPipeline} assigns (when {@code + * -DbeamTestPipelineUniqueRootNames=true}) before the test applies any transform, so that every + * step, PCollection and metric of a member is named {@code /...}. The coordinator never + * renames or otherwise mutates a member pipeline; a pipeline without a root name simply runs as its + * own job. + * + *

Threading model: a test thread only enqueues its pipeline and then blocks on a per-pipeline + * future. Per {@linkplain #compatibilityKey options-compatibility group}, a daemon scheduler thread + * drains the queue and groups pipelines into batches; each closed batch is handed to a daemon + * executor that submits the merged job, waits for it and reaches a verdict for every member. A + * member is credited from the merged job when its scoped PAsserts are verified and — if the + * job did not end {@code DONE} — every stage that may belong to it is {@code DONE} (see + * {@link StageAttribution}). Any other member gets {@link TestPipeline.StandaloneRerunRequested} + * thrown out of its {@code run()}, and its test is executed again from scratch as a standalone job + * by {@link BeamParallelJunit4Runner}. (A pipeline for which no merged job was built at all, + * because no partner arrived in time, simply runs standalone on its own test thread.) No test + * thread ever performs work on behalf of another test, so a test timeout or interrupt affects only + * that test. + * + *

No-hang guarantee: a member's future is always completed, whatever happens on the scheduler or + * batch-runner threads (see {@link BatchCollector} and {@link #executeBatch}); the merged job is + * waited on for at most the shortest member test timeout ({@link #mergedJobTimeout}); and as a last + * resort each test thread's wait is bounded ({@link PendingItem#awaitResult()}), so an unforeseen + * failure surfaces as a diagnosable test failure rather than a fork that never finishes. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +class DataflowTestBatchCoordinator { + private static final Logger LOG = LoggerFactory.getLogger(DataflowTestBatchCoordinator.class); + + static final String TEST_BATCHING_PROPERTY = "beam.dataflow.testBatching"; + static final String TEST_BATCH_MAX_SIZE_PROPERTY = "beam.dataflow.testBatchMaxSize"; + static final String TEST_BATCH_WINDOW_MS_PROPERTY = "beam.dataflow.testBatchWindowMs"; + + private static final long DEFAULT_BATCH_WINDOW_MS = 3000L; + private static final int DEFAULT_BATCH_MAX_SIZE = 30; + + /** + * Run-wide batching configuration. Batching spans many pipelines, so it is configured once per + * test JVM through system properties (see the ValidatesRunner tasks in {@code build.gradle}) + * rather than through any individual pipeline's {@link PipelineOptions}. + */ + @VisibleForTesting + static final class BatchingConfig { + final boolean enabled; + final int maxBatchSize; + final long windowMs; + + BatchingConfig(boolean enabled, int maxBatchSize, long windowMs) { + this.enabled = enabled; + this.maxBatchSize = maxBatchSize; + this.windowMs = windowMs; + } + + static BatchingConfig fromSystemProperties() { + // Unless sized explicitly, a batch may hold one pipeline per concurrently running test. + int defaultSize = + Integer.getInteger( + BeamParallelJunit4Runner.VALIDATES_RUNNER_THREADS_PROPERTY, DEFAULT_BATCH_MAX_SIZE); + return new BatchingConfig( + Boolean.getBoolean(TEST_BATCHING_PROPERTY), + Integer.getInteger(TEST_BATCH_MAX_SIZE_PROPERTY, Math.max(1, defaultSize)), + Long.getLong(TEST_BATCH_WINDOW_MS_PROPERTY, DEFAULT_BATCH_WINDOW_MS)); + } + } + + private static final AtomicInteger INSTANCE_COUNTER = new AtomicInteger(1); + private static volatile @Nullable DataflowTestBatchCoordinator shared; + + /** The process-wide coordinator shared by all {@link TestDataflowRunner} instances. */ + static DataflowTestBatchCoordinator shared() { + DataflowTestBatchCoordinator result = shared; + if (result == null) { + synchronized (DataflowTestBatchCoordinator.class) { + result = shared; + if (result == null) { + result = new DataflowTestBatchCoordinator(BatchingConfig.fromSystemProperties()); + shared = result; + } + } + } + return result; + } + + private final BatchingConfig config; + private final String name; + private final AtomicInteger batchCounter = new AtomicInteger(1); + private final AtomicInteger collectorCounter = new AtomicInteger(1); + private final ConcurrentHashMap collectors = new ConcurrentHashMap<>(); + private volatile boolean closed; + + /** Runs merged jobs so that schedulers keep collecting while jobs are in flight. */ + private final ExecutorService batchRunners; + + DataflowTestBatchCoordinator(BatchingConfig config) { + this.config = config; + this.name = "beam-test-batch-" + INSTANCE_COUNTER.getAndIncrement(); + this.batchRunners = Executors.newCachedThreadPool(daemonThreadFactory(name + "-runner-")); + } + + /** + * Stops the scheduler threads and the batch-runner pool. Members still waiting for a verdict fail + * promptly and later submissions are rejected. Tests only. + */ + @VisibleForTesting + void shutdown() { + closed = true; + collectors.values().forEach(BatchCollector::stop); + batchRunners.shutdownNow(); + } + + // ObjectMapper used to serialize PipelineOptions to determine batch keys. + private static final ObjectMapper MAPPER = + new ObjectMapper() + .registerModules(ObjectMapper.findModules(ReflectHelpers.findClassLoader())); + + /** + * Options that are expected to differ between tests which may nevertheless share a merged job. + * + *

The merged job is submitted with the first member's options only, so every other + * option that has been set on a member must be identical across the batch. Keep this list to + * per-test identity and local harness knobs; never add anything that influences how the service + * runs the job. + */ + @VisibleForTesting + static final ImmutableSet PER_TEST_OPTIONS = + ImmutableSet.of( + // Identity of the individual test / job. + "jobName", + "appName", + "optionsId", + // Derived per job from tempRoot + jobName by TestDataflowRunner.fromOptions. + "tempLocation", + "gcpTempLocation", + // Local harness behaviour; does not affect the submitted job. + "testTimeoutSeconds"); + + /** + * Returns a key identifying the group of pipelines that may share a merged job with one built + * from {@code options}, or {@code null} if the options cannot be serialized, in which case the + * pipeline must run standalone. + * + *

The key is a digest of every option that has been explicitly set on the options object + * through any {@link PipelineOptions} interface it has been viewed as, minus {@link + * #PER_TEST_OPTIONS}. Any option added to any interface (worker pool, debug, GCP, SDK harness, + * user-defined, ...) automatically splits batches when it differs instead of being silently + * dropped for all but the first member. + * + *

Defaults that were merely bound by a getter are left out. The serialized form contains them + * too, but which defaults are bound depends on which getters the test's transforms happened to + * call during construction (for example {@code View.asList()} and {@code Reshuffle} read {@code + * updateCompatibilityVersion}), and two tests must not end up in different batches because one of + * them read an option both would have received the same default for. {@link + * org.apache.beam.sdk.testing.TestPipeline} hands the runner a JSON copy of the test's options; + * it serializes that copy with {@link PipelineOptionsFactory#SERIALIZE_DEFAULTS_ATTRIBUTE} so + * that the distinction survives the round trip. + */ + @VisibleForTesting + static @Nullable String compatibilityKey(PipelineOptions options) { + String canonical = canonicalOptionsJson(options); + return canonical == null ? null : digest(canonical); + } + + private static String digest(String canonicalOptions) { + return Hashing.sha256().hashString(canonicalOptions, StandardCharsets.UTF_8).toString(); + } + + @VisibleForTesting + static @Nullable String canonicalOptionsJson(PipelineOptions options) { + JsonNode serialized; + Set explicitlySet; + try { + serialized = MAPPER.valueToTree(options).get("options"); + explicitlySet = PipelineOptionsFactory.explicitlySetProperties(options); + } catch (RuntimeException e) { + LOG.warn( + "Pipeline options for {} cannot be serialized; the pipeline will run standalone instead" + + " of in a merged Dataflow test job.", + options.as(ApplicationNameOptions.class).getAppName(), + e); + return null; + } + if (serialized == null || !serialized.isObject()) { + return null; + } + ObjectNode filtered = MAPPER.createObjectNode(); + for (Map.Entry entry : sortedFields(serialized).entrySet()) { + if (explicitlySet.contains(entry.getKey()) && !PER_TEST_OPTIONS.contains(entry.getKey())) { + filtered.set(entry.getKey(), canonicalize(entry.getValue())); + } + } + return filtered.toString(); + } + + /** Recursively orders object fields so that structurally equal values serialize identically. */ + private static JsonNode canonicalize(JsonNode node) { + if (node.isObject()) { + ObjectNode out = MAPPER.createObjectNode(); + for (Map.Entry entry : sortedFields(node).entrySet()) { + out.set(entry.getKey(), canonicalize(entry.getValue())); + } + return out; + } + if (node.isArray()) { + ArrayNode out = MAPPER.createArrayNode(); + for (JsonNode element : node) { + out.add(canonicalize(element)); + } + return out; + } + return node; + } + + private static TreeMap sortedFields(JsonNode node) { + TreeMap sorted = new TreeMap<>(); + node.fields().forEachRemaining(entry -> sorted.put(entry.getKey(), entry.getValue())); + return sorted; + } + + boolean isEligibleForBatching(Pipeline pipeline, TestDataflowPipelineOptions options) { + if (!config.enabled || config.maxBatchSize <= 1) { + return false; + } + if (!options.isBlockOnRun()) { + return false; + } + if (options.isStreaming()) { + // Streaming batches need their own verdict strategy (see executeBatchInternal); until that + // exists streaming tests always run as individual jobs. Say so once so the flag is not + // silently ignored by the streaming ValidatesRunner tasks. + if (STREAMING_UNSUPPORTED_LOGGED.compareAndSet(false, true)) { + LOG.info( + "Dataflow test batching does not yet support streaming jobs; streaming tests run as" + + " individual jobs."); + } + return false; + } + if (!(pipeline instanceof TestPipeline)) { + return false; + } + TestPipeline testPipeline = (TestPipeline) pipeline; + if (testPipeline.isStandaloneExecutionRequired()) { + return false; + } + if (testPipeline.getRootName().isEmpty()) { + // Members of a merged job are told apart purely by their root name, so an unnamed pipeline + // can never be merged. + if (UNNAMED_PIPELINE_WARNED.compareAndSet(false, true)) { + LOG.warn( + "Dataflow test batching is enabled but test pipelines have no unique root name, so" + + " every test runs as its own job. Set -D{}=true to enable merging.", + TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES); + } + return false; + } + if (!isDefaultMatcher(options.getOnCreateMatcher()) + || !isDefaultMatcher(options.getOnSuccessMatcher())) { + return false; + } + EligibilityVisitor visitor = new EligibilityVisitor(); + pipeline.traverseTopologically(visitor); + return visitor.hasPrimitiveTransform && !visitor.hasUnboundedPCollection; + } + + private static final AtomicBoolean UNNAMED_PIPELINE_WARNED = new AtomicBoolean(false); + private static final AtomicBoolean STREAMING_UNSUPPORTED_LOGGED = new AtomicBoolean(false); + + private static boolean isDefaultMatcher(@Nullable Object matcher) { + return matcher == null + || matcher instanceof TestPipelineOptions.AlwaysPassMatcher + || SerializableMatchers.anything().equals(matcher); + } + + /** + * Submits {@code pipeline} for batched execution and blocks until a verdict has been reached. + * Callers must have checked {@link #isEligibleForBatching} first; in particular the pipeline must + * have a unique root name. + */ + DataflowPipelineJob runInBatch( + TestPipeline pipeline, + TestDataflowPipelineOptions options, + TestDataflowRunner runner, + DataflowRunner delegateRunner) { + final String canonicalOptions = canonicalOptionsJson(options); + if (canonicalOptions == null) { + return runner.runStandalone(pipeline, delegateRunner); + } + String key = digest(canonicalOptions); + PendingItem item = + new PendingItem(pipeline, options, runner, delegateRunner, verdictTimeoutMillis(options)); + while (true) { + if (closed) { + throw new IllegalStateException( + "Dataflow test batch coordinator " + name + " has been shut down"); + } + BatchCollector collector = + collectors.computeIfAbsent( + key, + k -> { + BatchCollector created = + new BatchCollector(k, name + "-group-" + collectorCounter.getAndIncrement()); + LOG.info( + "Created Dataflow test batch group {} ({}) for options compatible with {}", + created.name, + k, + options.getAppName()); + LOG.debug("Batch group {} options: {}", created.name, canonicalOptions); + created.start(); + return created; + }); + if (collector.enqueue(item)) { + // A collector created after shutdown() iterated the registry would never be stopped and + // the item would sit until its verdict timeout. Stopping is idempotent, and close() fails + // whatever is queued, so the item gets a prompt verdict either way. + if (closed) { + collector.stop(); + } + break; + } + // Collectors never close on their own: enqueue only fails if this coordinator is shutting + // down (the re-check at the top of the loop then throws) or the scheduler thread died from an + // Error, in which case a fresh collector takes over for later tests. close() unregisters the + // collector too, but sets `closed` first; removing here as well covers that window so the + // loop cannot spin on the same dead instance. + collectors.remove(key, collector); + } + return item.awaitResult(); + } + + /** + * How long a test thread waits for the batch machinery's verdict before failing its test, or a + * negative value for no limit. Every path through the machinery completes the member's future, so + * this is a safety net that turns an unknown bug into a diagnosable failure rather than a hung + * fork. It is deliberately generous: the merged job itself is only waited on for the shortest + * member test timeout (see {@link #mergedJobTimeout}); the rest is slack for queueing (at most + * two batch windows), submission, metrics retrieval and cancellation. + */ + private long verdictTimeoutMillis(TestDataflowPipelineOptions options) { + Long testTimeoutSeconds = options.getTestTimeoutSeconds(); + if (testTimeoutSeconds == null || testTimeoutSeconds <= 0) { + return -1; + } + return 2 * config.windowMs + 2 * TimeUnit.SECONDS.toMillis(testTimeoutSeconds); + } + + /** + * The bound on waiting for a merged job: the shortest {@code testTimeoutSeconds} among its + * members, so that no member waits longer for the merged job than it would have for its own + * standalone job. Negative (unbounded) only if no member has a timeout configured. + */ + @VisibleForTesting + static Duration mergedJobTimeout(List items) { + long minSeconds = Long.MAX_VALUE; + for (PendingItem item : items) { + Long timeout = item.options.getTestTimeoutSeconds(); + if (timeout != null && timeout > 0) { + minSeconds = Math.min(minSeconds, timeout); + } + } + return Duration.standardSeconds(minSeconds == Long.MAX_VALUE ? -1 : minSeconds); + } + + private static class EligibilityVisitor extends PipelineVisitor.Defaults { + private boolean hasPrimitiveTransform = false; + private boolean hasUnboundedPCollection = false; + + @Override + public void visitPrimitiveTransform(TransformHierarchy.Node node) { + hasPrimitiveTransform = true; + } + + @Override + public void visitValue(PValue value, TransformHierarchy.Node producer) { + if (value instanceof PCollection + && ((PCollection) value).isBounded() == IsBounded.UNBOUNDED) { + hasUnboundedPCollection = true; + } + } + } + + /** + * The verdict the batch machinery reaches for one member pipeline: credited with a {@link #job} + * from the merged run, declined before any merged job was built from it, or in need of a fresh + * standalone execution after a merged attempt that yielded no verdict for it. + */ + @VisibleForTesting + static final class BatchOutcome { + enum Kind { + /** Credited from the merged job; {@link #job} is the member's scoped view of it. */ + PASSED, + /** + * Batching was declined before any merged job was built from the pipeline (for example, no + * partner arrived within the window). The pipeline is untouched, so the member's own test + * thread runs it standalone exactly as if it had been ineligible. + */ + DECLINED, + /** + * A merged job was built from the pipeline but did not yield a verdict for this member. The + * pipeline object is not run again (building the merged job rewrote it with another runner); + * the member's test is re-executed from scratch instead, see {@link + * TestPipeline.StandaloneRerunRequested}. + */ + RERUN_REQUIRED + } + + final Kind kind; + final @Nullable DataflowPipelineJob job; + + /** Why no verdict could be reached; {@link Kind#RERUN_REQUIRED} only. */ + final @Nullable String reason; + + private BatchOutcome(Kind kind, @Nullable DataflowPipelineJob job, @Nullable String reason) { + this.kind = kind; + this.job = job; + this.reason = reason; + } + + static BatchOutcome passed(DataflowPipelineJob job) { + return new BatchOutcome(Kind.PASSED, job, null); + } + + static BatchOutcome declined() { + return new BatchOutcome(Kind.DECLINED, null, null); + } + + static BatchOutcome rerunRequired(String reason) { + return new BatchOutcome(Kind.RERUN_REQUIRED, null, reason); + } + } + + @VisibleForTesting + static final class PendingItem { + final TestPipeline pipeline; + final TestDataflowPipelineOptions options; + final TestDataflowRunner runner; + final DataflowRunner delegateRunner; + final int expectedAssertions; + final long verdictTimeoutMillis; + final CompletableFuture resultFuture = new CompletableFuture<>(); + + /** + * The pipeline's unique root name (see {@link TestPipeline#getRootName()}). Every transform and + * PCollection of this member is named {@code /...}, which is how the member's steps, + * stages and metrics are told apart from the other members' in the merged job. + */ + final String scope; + + /** Assigned by the batch runner once the merged job is submitted; {@code null} before that. */ + volatile @Nullable String mergedJobId; + + PendingItem( + TestPipeline pipeline, + TestDataflowPipelineOptions options, + TestDataflowRunner runner, + DataflowRunner delegateRunner, + long verdictTimeoutMillis) { + this.pipeline = pipeline; + this.options = options; + this.runner = runner; + this.delegateRunner = delegateRunner; + this.verdictTimeoutMillis = verdictTimeoutMillis; + this.expectedAssertions = PAssert.countAsserts(pipeline); + this.scope = pipeline.getRootName(); + checkArgument(!scope.isEmpty(), "Batched test pipelines must have a root name"); + } + + /** + * Blocks the calling (test) thread until the batch machinery has reached a verdict, then + * returns the merged job, runs this pipeline standalone on the calling thread (if batching was + * declined before a merged job was built from it), or throws {@link + * TestPipeline.StandaloneRerunRequested} so that the test runner executes the test again from + * scratch as a job of its own. + * + *

The wait is bounded by {@link #verdictTimeoutMillis}. Exceeding it means the machinery + * failed to deliver a verdict at all, which is a harness bug; the test fails with a description + * of where the member got stuck rather than silently hanging its fork. + */ + DataflowPipelineJob awaitResult() { + BatchOutcome outcome; + try { + outcome = + verdictTimeoutMillis < 0 + ? resultFuture.get() + : resultFuture.get(verdictTimeoutMillis, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException( + "Interrupted while waiting for the merged Dataflow test batch containing " + + options.getAppName(), + e); + } catch (TimeoutException e) { + String jobId = mergedJobId; + throw new IllegalStateException( + String.format( + "No verdict for test %s from the Dataflow test batch machinery after %d ms" + + " (batch scope: %s; merged job: %s). This is a bug in the test batching" + + " machinery, not in the test.", + options.getAppName(), + verdictTimeoutMillis, + scope, + jobId == null ? "never submitted" : jobId), + e); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } else if (cause instanceof Error) { + throw (Error) cause; + } + throw new RuntimeException(cause != null ? cause : e); + } + switch (outcome.kind) { + case PASSED: + return outcome.job; + case DECLINED: + return runner.runStandalone(pipeline, delegateRunner); + case RERUN_REQUIRED: + default: + String jobId = mergedJobId; + throw new TestPipeline.StandaloneRerunRequested( + String.format( + "Test %s (scope %s) was merged into Dataflow job %s, which yielded no verdict" + + " for it: %s. The test has to be executed again as a job of its own;" + + " BeamParallelJunit4Runner does so automatically.", + options.getAppName(), + scope, + jobId == null ? "(never submitted)" : jobId, + outcome.reason)); + } + } + } + + private static ThreadFactory daemonThreadFactory(String namePrefix) { + AtomicInteger counter = new AtomicInteger(1); + return runnable -> { + Thread thread = new Thread(runnable, namePrefix + counter.getAndIncrement()); + thread.setDaemon(true); + return thread; + }; + } + + /** + * Groups pipelines that share an options-compatibility key into batches. + * + *

Test threads only {@link #enqueue} and then wait on their {@link PendingItem#resultFuture}. + * A single daemon scheduler thread drains the queue: it takes the first pipeline, waits up to the + * batch window for more (closing early once the batch is full), and hands the closed batch to the + * coordinator's {@link #batchRunners} so that it can immediately start collecting the next batch + * while the merged job is in flight. A batch of one is not worth a merged job and is sent + * straight back to its test thread to run standalone. + * + *

Lifecycle: however the scheduler thread ends (interrupt, or an unexpected error), the + * collector {@link #close}s: it marks itself closed, unregisters from the coordinator so that the + * next submission gets a fresh collector, and fails everything still queued. {@link #enqueue} + * cooperates with that so no item can be added after the final drain and never be seen again. + */ + private final class BatchCollector { + private final String key; + private final String name; + private final LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + private @Nullable Thread scheduler; + + /** Once set, nothing in {@link #queue} will ever be dispatched; see {@link #enqueue}. */ + private volatile boolean closed; + + BatchCollector(String key, String name) { + this.key = key; + this.name = name; + } + + synchronized void start() { + if (scheduler == null) { + scheduler = new Thread(this::schedulerLoop, name + "-scheduler"); + scheduler.setDaemon(true); + scheduler.start(); + } + } + + synchronized void stop() { + if (scheduler != null) { + scheduler.interrupt(); + } + } + + /** + * Returns {@code false} if this collector is already closed and the caller must use another + * one. Returns {@code true} once {@code item} is guaranteed a verdict: either the scheduler + * dispatches it, or it is failed when the collector closes. + */ + boolean enqueue(PendingItem item) { + if (closed) { + return false; + } + queue.add(item); + if (closed) { + // Raced with close(): its drain may have run before our add. Draining again is idempotent + // (futures complete at most once), so whichever side sees the item fails it. + failQueued(null); + } + return true; + } + + private void schedulerLoop() { + Throwable failure = null; + try { + while (true) { + dispatch(collectBatch()); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Throwable t) { + failure = t; + LOG.error( + "Dataflow test batch scheduler {} died unexpectedly; queued tests will fail and later" + + " tests will use a new scheduler.", + name, + t); + } finally { + close(failure); + } + } + + private void close(@Nullable Throwable cause) { + closed = true; + collectors.remove(key, this); + failQueued(cause); + } + + private void failQueued(@Nullable Throwable cause) { + List orphans = new ArrayList<>(); + queue.drainTo(orphans); + for (PendingItem orphan : orphans) { + orphan.resultFuture.completeExceptionally( + new IllegalStateException( + "Dataflow test batch scheduler " + + name + + " stopped before scheduling " + + orphan.options.getAppName(), + cause)); + } + } + + /** + * Blocks until a member arrives, then keeps collecting for up to the batch window or until the + * batch is full. Never returns an empty batch. + */ + private List collectBatch() throws InterruptedException { + List batch = new ArrayList<>(); + try { + batch.add(queue.take()); + // The window opens when the first member arrives, not while waiting for it. + long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(config.windowMs); + while (batch.size() < config.maxBatchSize) { + long remainingNanos = deadlineNanos - System.nanoTime(); + if (remainingNanos <= 0) { + break; + } + PendingItem next = queue.poll(remainingNanos, TimeUnit.NANOSECONDS); + if (next == null) { + break; + } + batch.add(next); + } + } catch (InterruptedException e) { + // Hand the partial batch back so that close() fails these members too. + queue.addAll(batch); + throw e; + } + return batch; + } + + /** + * Hands a closed batch off without doing any job work on the scheduler thread; this returns + * immediately either way. + * + *

A batch of one is not executed through {@link #executeBatch}: that would wrap the lone + * pipeline in a renamed composite, apply scoped PAssert/stage verification, and — on + * failure — require the test to be executed again. Instead the member is told that + * batching was declined, which makes it behave exactly like an ineligible pipeline: its own + * test thread runs it standalone (see {@link PendingItem#awaitResult()}) with the full {@link + * TestDataflowRunner} semantics and unmodified names. + */ + private void dispatch(List batch) { + try { + if (batch.size() == 1) { + // Verdict only; the job is submitted by the test thread. + batch.get(0).resultFuture.complete(BatchOutcome.declined()); + return; + } + batchRunners.execute(() -> executeBatch(batch)); + } catch (Throwable t) { + for (PendingItem item : batch) { + item.resultFuture.completeExceptionally(t); + } + } + } + } + + /** + * Runs one merged Dataflow job for {@code items} and completes every member's future. Never + * throws and never leaves a future incomplete: an unexpected failure anywhere in here is + * propagated to all member test threads instead of hanging them. + */ + @VisibleForTesting + void executeBatch(List items) { + try { + executeBatchInternal(items); + } catch (Throwable t) { + LOG.error( + "Unexpected failure while executing merged Dataflow test batch; failing {} tests.", + items.size(), + t); + for (PendingItem item : items) { + item.resultFuture.completeExceptionally(t); + } + } finally { + for (PendingItem item : items) { + if (!item.resultFuture.isDone()) { + item.resultFuture.completeExceptionally( + new IllegalStateException( + "Merged Dataflow test batch finished without a verdict for " + + item.options.getAppName())); + } + } + } + } + + /** + * The most conservative Dataflow job-name length we know of. The current REST reference allows + * {@code [a-z]([-a-z0-9]{0,1022}[a-z0-9])?}, but the service historically enforced 63 (and the + * Python SDK still tests against that), so merged job names are kept within it. + */ + @VisibleForTesting static final int MAX_JOB_NAME_LENGTH = 63; + + /** + * Every merged job's name starts with this, which is how a pipeline running on Dataflow can tell + * (via its {@code jobName} option) that it is part of a merged job rather than a standalone one. + */ + static final String BATCH_JOB_NAME_PREFIX = "batch-"; + + private static final DateTimeFormatter JOB_NAME_TIMESTAMP = + DateTimeFormat.forPattern("MMddHHmmss").withZone(DateTimeZone.UTC); + + /** Builds the merged job's name; see {@link #batchJobName(int, String, long, int)}. */ + private static String batchJobName(int batchNumber, @Nullable String leaderAppName) { + return batchJobName( + batchNumber, + leaderAppName, + System.currentTimeMillis(), + ThreadLocalRandom.current().nextInt()); + } + + /** + * Builds {@code batch----} within {@link + * #MAX_JOB_NAME_LENGTH}. + * + *

Uniqueness comes from the timestamp and random suffix, the same scheme as {@code + * PipelineOptions.JobNameFactory} uses for standalone jobs. It must not come from the batch + * number: that counter is per JVM, and the same test classes run concurrently in several JVMs and + * postcommits that all count from one. A colliding name makes Dataflow return the other job and + * the whole batch falls back to standalone runs. The leader's app name is only a hint for humans + * reading the console and is the only part that is ever truncated. + */ + @VisibleForTesting + static String batchJobName( + int batchNumber, @Nullable String leaderAppName, long nowMillis, int random) { + String prefix = BATCH_JOB_NAME_PREFIX + batchNumber; + String suffix = "-" + JOB_NAME_TIMESTAMP.print(nowMillis) + "-" + Integer.toHexString(random); + String hint = + leaderAppName == null ? "" : leaderAppName.toLowerCase().replaceAll("[^a-z0-9]", "0"); + int room = MAX_JOB_NAME_LENGTH - prefix.length() - suffix.length() - 1; + if (hint.length() > room) { + hint = hint.substring(0, Math.max(0, room)); + } + return hint.isEmpty() ? prefix + suffix : prefix + "-" + hint + suffix; + } + + private void executeBatchInternal(List items) { + int batchNumber = batchCounter.getAndIncrement(); + List memberPipelines = new ArrayList<>(items.size()); + List scopes = new ArrayList<>(items.size()); + StringBuilder membership = new StringBuilder(); + for (int i = 0; i < items.size(); i++) { + PendingItem item = items.get(i); + memberPipelines.add(item.pipeline); + scopes.add(item.scope); + membership + .append(i == 0 ? "" : ", ") + .append(item.scope) + .append('=') + .append(item.options.getAppName()); + } + LOG.info( + "Merging {} ValidatesRunner test pipelines into Dataflow batch #{}: {}", + items.size(), + batchNumber, + membership); + + PendingItem leader = items.get(0); + String originalLeaderJobName = leader.options.getJobName(); + String batchJobName = batchJobName(batchNumber, leader.options.getAppName()); + leader.options.setJobName(batchJobName); + + DataflowPipelineJob batchJob = null; + State terminalState = null; + ErrorMonitorMessagesHandler messages = null; + String noVerdictReason = null; + try { + // The constructor rejects members whose transform names collide. That cannot happen with + // TestPipeline's sequence-numbered root names, but if it did the catch below sends every + // member back for a fresh standalone execution rather than risk mis-crediting a test. + CompositePipeline compositePipeline = new CompositePipeline(leader.options, memberPipelines); + batchJob = leader.runner.submit(leader.delegateRunner, compositePipeline); + for (PendingItem item : items) { + item.mergedJobId = batchJob.getJobId(); + } + LOG.info( + "Submitted merged Dataflow job {} for batch #{} ({} tests)", + batchJob.getJobId(), + batchNumber, + items.size()); + messages = TestDataflowRunner.errorMonitorFor(batchJob); + Duration timeout = mergedJobTimeout(items); + terminalState = leader.runner.waitForMergedJobTermination(batchJob, timeout, messages); + if (terminalState == null) { + noVerdictReason = + "the merged job did not terminate within " + timeout + " and was cancelled"; + LOG.warn( + "Merged Dataflow job {} for batch #{} did not terminate within {}; cancelling it; all" + + " {} tests will be executed again standalone.{}", + batchJob.getJobId(), + batchNumber, + timeout, + items.size(), + errorSuffix(messages)); + cancelQuietly(batchJob); + } + } catch (Exception e) { + noVerdictReason = "the merged job failed during submission or execution (" + e + ")"; + LOG.warn( + "Merged Dataflow batch #{} failed during submission or execution; all {} tests will be" + + " executed again standalone.{}", + batchNumber, + items.size(), + errorSuffix(messages), + e); + cancelQuietly(batchJob); + } finally { + leader.options.setJobName(originalLeaderJobName); + } + + List outcomes; + if (batchJob == null || terminalState == null) { + outcomes = new ArrayList<>(items.size()); + for (int i = 0; i < items.size(); i++) { + outcomes.add(BatchOutcome.rerunRequired(noVerdictReason + errorSuffix(messages))); + } + } else { + // How members are credited depends on the execution mode, and all members of a batch share + // one mode because `streaming` is part of the compatibility key. Only batch mode is eligible + // today (see isEligibleForBatching); a streaming strategy would plug in here, e.g. crediting + // everyone on DONE without error messages and nobody otherwise, since streaming jobs expose + // neither per-stage terminal states nor the batch verification counters. + checkState(!leader.options.isStreaming(), "Streaming batches are not supported yet"); + outcomes = batchModeVerdicts(batchJob, terminalState, messages, items, scopes); + } + + // Complete futures only after every member has been evaluated. + for (int i = 0; i < items.size(); i++) { + items.get(i).resultFuture.complete(outcomes.get(i)); + } + } + + /** + * Credits members of a terminated batch-mode merged job. A member is credited when its + * scoped PAssert counters verify and, unless the job ended {@code DONE}, every stage that may + * belong to it ended {@code DONE} (see {@link StageAttribution}); anyone else must be executed + * again standalone. + */ + private List batchModeVerdicts( + DataflowPipelineJob batchJob, + State terminalState, + ErrorMonitorMessagesHandler messages, + List items, + List scopes) { + List outcomes = new ArrayList<>(items.size()); + // Metrics are fetched regardless of terminal state: counters reported by stages that + // completed before the job failed are retained by the service. + JobMetrics metrics = leader(items).runner.getJobMetrics(batchJob); + StageAttribution attribution = null; + if (terminalState != State.DONE) { + attribution = + StageAttribution.fromJob(leader(items).runner.getJobWithExecutionDetails(batchJob)); + if (attribution == null) { + LOG.warn( + "Merged Dataflow job {} terminated in state {} and per-stage execution state is" + + " unavailable; all {} tests will be executed again standalone.{}", + batchJob.getJobId(), + terminalState, + items.size(), + errorSuffix(messages)); + } + } + for (PendingItem item : items) { + String scope = item.scope; + PAssertCounts counts = scopedPAssertCounts(batchJob, metrics, scope); + boolean countersOk = counts != null && counts.satisfy(item.expectedAssertions); + String passertVerdict = + (countersOk ? "PAsserts verified (" : "PAsserts unverified (") + + (counts == null ? "metrics unavailable" : counts) + + " out of " + + item.expectedAssertions + + " expected assertions)"; + String verdict; + boolean credited; + if (terminalState == State.DONE) { + credited = countersOk; + verdict = "job DONE, " + passertVerdict; + } else if (attribution == null) { + credited = false; + verdict = "job " + terminalState + ", stage states unavailable"; + } else { + StageAttribution.ScopeStatus status = attribution.statusForScope(scope, scopes); + credited = countersOk && status.isComplete(); + verdict = "job " + terminalState + ", " + status + ", " + passertVerdict; + } + if (credited) { + LOG.info( + "Test {} (scope {}) passed in merged Dataflow job {} ({}).", + item.options.getAppName(), + scope, + batchJob.getJobId(), + verdict); + outcomes.add(BatchOutcome.passed(new ScopedDataflowPipelineJob(batchJob, scope))); + } else { + LOG.warn( + "Test {} (scope {}) could not be credited from merged Dataflow job {} ({}); it will be" + + " executed again standalone.{}", + item.options.getAppName(), + scope, + batchJob.getJobId(), + verdict, + errorSuffix(messages)); + outcomes.add(BatchOutcome.rerunRequired(verdict + errorSuffix(messages))); + } + } + return outcomes; + } + + private static PendingItem leader(List items) { + return items.get(0); + } + + /** The job's error messages, formatted for appending to a log line, or empty if none. */ + private static String errorSuffix(@Nullable ErrorMonitorMessagesHandler messages) { + if (messages == null || !messages.hasSeenError()) { + return ""; + } + return " Job errors: " + messages.getErrorMessage(); + } + + private static void cancelQuietly(@Nullable DataflowPipelineJob job) { + if (job == null) { + return; + } + try { + if (!job.getState().isTerminal()) { + LOG.info("Cancelling abandoned merged Dataflow job {}", job.getJobId()); + job.cancel(); + } + } catch (Exception e) { + LOG.warn("Failed to cancel merged Dataflow job {}", job.getJobId(), e); + } + } + + /** + * Per-stage execution state of a Dataflow job, attributed to batch member scopes. + * + *

Attribution is purely by the user-facing names a stage reports (component transforms and + * sources, input and output sources): a stage "may belong to" a scope if any of its names is + * under that scope's prefix. No assumption is made about how the service fuses steps into stages: + * a stage that reports no names, or a name under no known scope, is treated as possibly belonging + * to every scope. This makes the verdict conservative in exactly the cases where the service's + * view cannot be mapped back to a member. + */ + @VisibleForTesting + static final class StageAttribution { + private static final String STAGE_DONE = "JOB_STATE_DONE"; + + private final Map stateByStage; + private final Map> namesByStage; + + private StageAttribution( + Map stateByStage, Map> namesByStage) { + this.stateByStage = stateByStage; + this.namesByStage = namesByStage; + } + + /** Returns {@code null} if the job does not carry both stage states and stage descriptions. */ + static @Nullable StageAttribution fromJob(@Nullable Job job) { + if (job == null + || job.getStageStates() == null + || job.getStageStates().isEmpty() + || job.getPipelineDescription() == null + || job.getPipelineDescription().getExecutionPipelineStage() == null + || job.getPipelineDescription().getExecutionPipelineStage().isEmpty()) { + return null; + } + Map stateByStage = new HashMap<>(); + for (ExecutionStageState stageState : job.getStageStates()) { + if (stageState.getExecutionStageName() != null + && stageState.getExecutionStageState() != null) { + stateByStage.put(stageState.getExecutionStageName(), stageState.getExecutionStageState()); + } + } + Map> namesByStage = new LinkedHashMap<>(); + for (ExecutionStageSummary stage : job.getPipelineDescription().getExecutionPipelineStage()) { + if (stage.getName() == null) { + continue; + } + Set names = new HashSet<>(); + if (stage.getComponentTransform() != null) { + for (ComponentTransform t : stage.getComponentTransform()) { + addName(names, t.getUserName()); + addName(names, t.getOriginalTransform()); + } + } + if (stage.getComponentSource() != null) { + for (ComponentSource s : stage.getComponentSource()) { + addName(names, s.getUserName()); + addName(names, s.getOriginalTransformOrCollection()); + } + } + addSourceNames(names, stage.getInputSource()); + addSourceNames(names, stage.getOutputSource()); + namesByStage.put(stage.getName(), names); + } + return new StageAttribution(stateByStage, namesByStage); + } + + private static void addSourceNames(Set names, @Nullable List sources) { + if (sources == null) { + return; + } + for (StageSource source : sources) { + addName(names, source.getUserName()); + addName(names, source.getOriginalTransformOrCollection()); + } + } + + private static void addName(Set names, @Nullable String name) { + if (name != null && !name.isEmpty()) { + names.add(name); + } + } + + /** Summarises the stages that may belong to one scope. */ + static final class ScopeStatus { + private final int doneStages; + private final List incompleteStages; + + private ScopeStatus(int doneStages, List incompleteStages) { + this.doneStages = doneStages; + this.incompleteStages = incompleteStages; + } + + /** True iff at least one stage may belong to the scope and every such stage is DONE. */ + boolean isComplete() { + return doneStages > 0 && incompleteStages.isEmpty(); + } + + @Override + public String toString() { + return doneStages + + " stage(s) done" + + (incompleteStages.isEmpty() ? "" : ", not done: " + incompleteStages); + } + } + + ScopeStatus statusForScope(String scope, Collection allScopes) { + int done = 0; + List incomplete = new ArrayList<>(); + for (Map.Entry> stage : namesByStage.entrySet()) { + if (!mayBelongToScope(stage.getValue(), scope, allScopes)) { + continue; + } + String state = stateByStage.get(stage.getKey()); + if (STAGE_DONE.equals(state)) { + done++; + } else { + incomplete.add(stage.getKey() + "=" + (state == null ? "UNKNOWN" : state)); + } + } + return new ScopeStatus(done, incomplete); + } + + private static boolean mayBelongToScope( + Set names, String scope, Collection allScopes) { + if (names.isEmpty()) { + return true; + } + for (String name : names) { + if (matchesScopePrefix(name, scope)) { + return true; + } + boolean underSomeScope = false; + for (String other : allScopes) { + if (matchesScopePrefix(name, other)) { + underSomeScope = true; + break; + } + } + if (!underSomeScope) { + return true; + } + } + return false; + } + } + + /** + * Sums the {@link PAssert} counters of the steps under {@code scopePrefix} in a merged job, or + * returns {@code null} if the job's metrics are unavailable. + */ + static @Nullable PAssertCounts scopedPAssertCounts( + DataflowPipelineJob batchJob, @Nullable JobMetrics metrics, String scopePrefix) { + if (metrics == null || metrics.getMetrics() == null) { + return null; + } + return TestDataflowRunner.countPAssertCounters( + metrics, + internalStepName -> { + if (internalStepName.isEmpty()) { + return false; + } + String userStepName = resolveUserStepName(batchJob, internalStepName); + return userStepName != null && matchesScopePrefix(userStepName, scopePrefix); + }); + } + + static @Nullable String resolveUserStepName(DataflowPipelineJob job, String internalStepName) { + RunnerApi.@Nullable Pipeline pipelineProto = job.getPipelineProto(); + if (pipelineProto != null) { + RunnerApi.@Nullable PTransform transform = + pipelineProto.getComponents().getTransformsMap().get(internalStepName); + if (transform != null && !transform.getUniqueName().isEmpty()) { + return transform.getUniqueName(); + } + } + BiMap, String> stepNames = job.getTransformStepNames(); + if (stepNames != null) { + AppliedPTransform applied = stepNames.inverse().get(internalStepName); + if (applied != null) { + return applied.getFullName(); + } + } + return internalStepName; + } + + /** + * Whether a step, stage or metric name belongs to the member whose root is {@code scopePrefix}. + * + *

Beam full names are {@code /}-separated ({@code t3-Foo-bar/ParDo(Anon)}), but Runner V2 + * stage descriptions and some step names come back mangled with {@code -} in place of {@code /} + * and output suffixes appended ({@code t3-Foo-bar-ParDo-Anon--out0}), so both separators are + * accepted. A separator (or an exact match) is required so that {@code t1} does not claim {@code + * t10/...}. + */ + static boolean matchesScopePrefix(String stepName, String scopePrefix) { + return stepName.startsWith(scopePrefix + "/") + || stepName.startsWith(scopePrefix + "-") + || stepName.equals(scopePrefix); + } + + /** Inverse of {@link #matchesScopePrefix}: removes the root and its separator, if present. */ + static String stripScopePrefix(String stepName, String scopePrefix) { + if (stepName.startsWith(scopePrefix + "/") || stepName.startsWith(scopePrefix + "-")) { + return stepName.substring(scopePrefix.length() + 1); + } + if (stepName.equals(scopePrefix)) { + return ""; + } + return stepName; + } + + /** + * A {@link DataflowPipelineJob} view for a single member pipeline within a merged batch job that + * filters and unprefixes {@link MetricResults} to that member's transform scope. + * + *

Instances are only ever handed to members that were credited by the coordinator, + * i.e. every execution stage belonging to the member finished successfully and its PAssert + * counters were satisfied. The underlying batch job may nevertheless be in a non-{@code DONE} + * terminal state because of a different member, so this view always reports {@link + * State#DONE} rather than delegating to the batch job's state. + * + *

The view deliberately holds no {@link DataflowClient}: it must never act on the shared batch + * job on behalf of one member (cancelling or draining it would abort the other members' work), + * and its state is already known. Every public operation that would otherwise contact the service + * — waiting, cancelling, draining, polling state — is therefore overridden to answer + * locally, and metrics are served through the batch job's own client. + */ + static final class ScopedDataflowPipelineJob extends DataflowPipelineJob { + private final MetricResults scopedMetrics; + + ScopedDataflowPipelineJob(DataflowPipelineJob delegate, String scopePrefix) { + super( + /* dataflowClient= */ null, + delegate.getJobId(), + delegate.getDataflowOptions(), + delegate.getTransformStepNames(), + delegate.getPipelineProto()); + this.scopedMetrics = new ScopedMetricResults(delegate, scopePrefix); + } + + @Override + public State getState() { + return State.DONE; + } + + @Override + public @Nullable State waitUntilFinish() { + return State.DONE; + } + + @Override + public @Nullable State waitUntilFinish(Duration duration) { + return State.DONE; + } + + @Override + public @Nullable State waitUntilFinish( + Duration duration, MonitoringUtil.JobMessagesHandler messageHandler) { + return State.DONE; + } + + @Override + public State cancel() { + // The member's work is already complete; cancelling is a no-op and must not affect the + // other members sharing the batch job. + return State.DONE; + } + + @Override + public State drain() { + // As for cancel(): draining the shared batch job would abort the other members. + return State.DONE; + } + + @Override + public MetricResults metrics() { + return scopedMetrics; + } + } + + private static final class ScopedMetricResults extends MetricResults { + private final DataflowPipelineJob delegateJob; + private final String scopePrefix; + + private ScopedMetricResults(DataflowPipelineJob delegateJob, String scopePrefix) { + this.delegateJob = delegateJob; + this.scopePrefix = scopePrefix; + } + + @Override + public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) { + MetricsFilter.Builder nameOnlyFilterBuilder = MetricsFilter.builder(); + if (filter != null) { + for (MetricNameFilter nameFilter : filter.names()) { + nameOnlyFilterBuilder.addNameFilter(nameFilter); + } + } + MetricQueryResults rawResults; + synchronized (delegateJob) { + rawResults = delegateJob.metrics().queryMetrics(nameOnlyFilterBuilder.build()); + } + return MetricQueryResults.create( + filterAndStrip(rawResults.getCounters(), filter), + filterAndStrip(rawResults.getDistributions(), filter), + filterAndStrip(rawResults.getGauges(), filter), + filterAndStrip(rawResults.getStringSets(), filter), + filterAndStrip(rawResults.getBoundedTries(), filter), + Collections.>emptyList()); + } + + private List> filterAndStrip( + Iterable> results, @Nullable MetricsFilter filter) { + List> scoped = new ArrayList<>(); + for (MetricResult result : results) { + MetricKey key = result.getKey(); + String stepName = key.stepName(); + if (stepName == null || !matchesScopePrefix(stepName, scopePrefix)) { + continue; + } + String unprefixedStep = stripScopePrefix(stepName, scopePrefix); + MetricKey unprefixedKey = MetricKey.create(unprefixedStep, key.metricName()); + if (filter == null || MetricFiltering.matches(filter, unprefixedKey)) { + scoped.add( + MetricResult.create( + unprefixedKey, result.getCommittedOrNull(), result.getAttempted())); + } + } + return scoped; + } + } +} diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java index 89d198154ad3..07e809dedfc3 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java @@ -19,29 +19,41 @@ import static org.hamcrest.MatcherAssert.assertThat; +import com.google.api.client.googleapis.json.GoogleJsonResponseException; +import com.google.api.client.util.BackOff; +import com.google.api.client.util.BackOffUtils; +import com.google.api.client.util.Sleeper; +import com.google.api.services.dataflow.model.Job; import com.google.api.services.dataflow.model.JobMessage; import com.google.api.services.dataflow.model.JobMetrics; +import com.google.api.services.dataflow.model.MetricStructuredName; import com.google.api.services.dataflow.model.MetricUpdate; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.io.File; import java.io.IOException; import java.math.BigDecimal; import java.util.List; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.Semaphore; +import java.util.function.Predicate; import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; import org.apache.beam.runners.dataflow.util.MonitoringUtil; import org.apache.beam.runners.dataflow.util.MonitoringUtil.JobMessagesHandler; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult.State; import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.extensions.gcp.util.BackOffAdapter; import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testing.TestPipelineOptions; +import org.apache.beam.sdk.util.FluentBackoff; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Optional; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.slf4j.Logger; @@ -58,17 +70,66 @@ }) public class TestDataflowRunner extends PipelineRunner { private static final String TENTATIVE_COUNTER = "tentative"; + static final String MAX_CONCURRENT_STANDALONE_JOBS_PROPERTY = + "beam.dataflow.maxConcurrentStandaloneJobs"; + private static final int DEFAULT_MAX_CONCURRENT_STANDALONE_JOBS = 12; + + /** + * Per-JVM cap on concurrently running standalone jobs. The permit count is a system property set + * once per test JVM by Gradle ({@link #MAX_CONCURRENT_STANDALONE_JOBS_PROPERTY}). + * + *

Every path that runs a test as a job of its own takes a permit here: tests that are not + * eligible for batching, tests whose batch turned out to be a batch of one, and tests executed + * again standalone after a merged job yielded no verdict for them. The last group can arrive in a + * burst (every member of an abandoned merged job at once), which is exactly when a shared cap + * matters. Merged jobs do not take a permit: there are only ever a handful in flight per JVM + * (bounded by test threads divided by batch size), each stands in for up to a whole batch of + * tests, and queueing one behind single-test jobs would invert that priority. + * + *

Fair, so the test that has waited longest for a permit is next: with many test threads per + * JVM an unfair semaphore can let a test starve behind a steady stream of newer arrivals. + */ + private static final Semaphore STANDALONE_SEMAPHORE = + new Semaphore( + Math.max( + 1, + Integer.getInteger( + MAX_CONCURRENT_STANDALONE_JOBS_PROPERTY, DEFAULT_MAX_CONCURRENT_STANDALONE_JOBS)), + true); + private static final Logger LOG = LoggerFactory.getLogger(TestDataflowRunner.class); + /** + * Backoff for job submissions rejected by the service for quota (HTTP 429). Many test JVMs share + * the project's job-creation quota with no coordination between them, so a rejection usually just + * means a brief burst; waiting it out beats failing the test. Roughly 30s, 1m, 2m, 2m, 2m (with + * jitter), i.e. up to five retries over about eight minutes. + */ + private static final FluentBackoff SUBMISSION_BACKOFF_FACTORY = + FluentBackoff.DEFAULT + .withInitialBackoff(Duration.standardSeconds(30)) + .withMaxBackoff(Duration.standardMinutes(2)) + .withMaxRetries(5); + private final TestDataflowPipelineOptions options; private final DataflowClient dataflowClient; private final DataflowRunner runner; + private final DataflowTestBatchCoordinator batchCoordinator; + private Sleeper sleeper = Sleeper.DEFAULT; private int expectedNumberOfAssertions = 0; TestDataflowRunner(TestDataflowPipelineOptions options, DataflowClient client) { + this(options, client, DataflowTestBatchCoordinator.shared()); + } + + TestDataflowRunner( + TestDataflowPipelineOptions options, + DataflowClient client, + DataflowTestBatchCoordinator batchCoordinator) { this.options = options; this.dataflowClient = client; this.runner = DataflowRunner.fromOptions(options); + this.batchCoordinator = batchCoordinator; } /** Constructs a runner from the provided options. */ @@ -102,17 +163,115 @@ static TestDataflowRunner fromOptionsAndClient( return new TestDataflowRunner(options, client); } + @VisibleForTesting + static TestDataflowRunner fromOptionsAndClient( + TestDataflowPipelineOptions options, + DataflowClient client, + DataflowTestBatchCoordinator batchCoordinator) { + return new TestDataflowRunner(options, client, batchCoordinator); + } + @Override public DataflowPipelineJob run(Pipeline pipeline) { return run(pipeline, runner); } DataflowPipelineJob run(Pipeline pipeline, DataflowRunner runner) { + return batchCoordinator.isEligibleForBatching(pipeline, options) + ? batchCoordinator.runInBatch((TestPipeline) pipeline, options, this, runner) + : runStandalone(pipeline, runner); + } + + @VisibleForTesting + void setSleeper(Sleeper sleeper) { + this.sleeper = sleeper; + } + + /** + * Submits {@code pipeline} through {@code runner}, retrying with backoff while the service + * rejects the submission for quota (HTTP 429). Any other failure propagates immediately. + */ + DataflowPipelineJob submit(DataflowRunner runner, Pipeline pipeline) { + BackOff backOff = BackOffAdapter.toGcpBackOff(SUBMISSION_BACKOFF_FACTORY.backoff()); + while (true) { + try { + return runner.run(pipeline); + } catch (RuntimeException e) { + GoogleJsonResponseException quota = quotaRejection(e); + if (quota == null) { + throw e; + } + boolean retry; + try { + retry = BackOffUtils.next(sleeper, backOff); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw e; + } catch (IOException ioe) { + throw e; + } + if (!retry) { + LOG.warn("Dataflow job submission still rejected for quota after retries; giving up."); + throw e; + } + LOG.warn( + "Dataflow job submission rejected for quota ({}); retrying.", + quota.getDetails() != null ? quota.getDetails().getMessage() : quota.getMessage()); + } + } + } + + /** The HTTP 429 response behind {@code t}, if that is what caused it. */ + @VisibleForTesting + static @Nullable GoogleJsonResponseException quotaRejection(Throwable t) { + java.util.Optional handledCause = + Throwables.getCausalChain(t).stream() + .filter(cause -> cause instanceof GoogleJsonResponseException) + .map(cause -> (GoogleJsonResponseException) cause) + .filter(cause -> cause.getStatusCode() == 429) + .findAny(); + return handledCause.orElse(null); + } + + /** + * Runs {@code pipeline} as its own Dataflow job, holding a permit from {@link + * #STANDALONE_SEMAPHORE} for its duration. + */ + DataflowPipelineJob runStandalone(Pipeline pipeline, DataflowRunner runner) { + boolean acquired = false; + try { + STANDALONE_SEMAPHORE.acquire(); + acquired = true; + return runStandaloneInternal(pipeline, runner); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } finally { + if (acquired) { + STANDALONE_SEMAPHORE.release(); + } + } + } + + /** + * Fetches {@code job} with {@code JOB_VIEW_ALL}, which carries per-stage execution states and the + * stage-to-user-step description needed to attribute stages to batch members. Returns {@code + * null} if the service call fails. + */ + @Nullable Job getJobWithExecutionDetails(DataflowPipelineJob job) { + try { + return dataflowClient.getJob(job.getJobId(), "JOB_VIEW_ALL"); + } catch (IOException e) { + LOG.warn("Failed to get execution details for Dataflow job {}: ", job.getJobId(), e); + return null; + } + } + + private DataflowPipelineJob runStandaloneInternal(Pipeline pipeline, DataflowRunner runner) { updatePAssertCount(pipeline); TestPipelineOptions testPipelineOptions = options.as(TestPipelineOptions.class); - final DataflowPipelineJob job; - job = runner.run(pipeline); + final DataflowPipelineJob job = submit(runner, pipeline); LOG.info( "Running Dataflow job {} with {} expected assertions.", @@ -203,6 +362,32 @@ private boolean waitForStreamingJobTermination( } } + /** + * Waits up to {@code timeout} for a merged test job to terminate, feeding its log messages to + * {@code messageHandler}. Returns its terminal state, or {@code null} if the job did not + * terminate in time or the wait was interrupted. + * + *

Unlike a standalone batch job, which this runner waits on indefinitely, a merged job holds + * the verdict for several tests at once, so it is never allowed to block them forever. + */ + @Nullable State waitForMergedJobTermination( + DataflowPipelineJob job, Duration timeout, ErrorMonitorMessagesHandler messageHandler) { + try { + State state = job.waitUntilFinish(timeout, messageHandler); + return state != null && state.isTerminal() ? state : null; + } catch (IOException e) { + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return null; + } + } + + /** An {@link ErrorMonitorMessagesHandler} for {@code job} that also logs every message. */ + static ErrorMonitorMessagesHandler errorMonitorFor(DataflowPipelineJob job) { + return new ErrorMonitorMessagesHandler(job, new MonitoringUtil.LoggingHandler()); + } + /** Return {@code true} if job state is {@code State.DONE}. {@code false} otherwise. */ private boolean waitForBatchJobTermination( DataflowPipelineJob job, ErrorMonitorMessagesHandler messageHandler) { @@ -259,21 +444,9 @@ Optional checkForPAssertSuccess(DataflowPipelineJob job) { return Optional.absent(); } - int successes = 0; - int failures = 0; - for (MetricUpdate metric : metrics.getMetrics()) { - if (metric.getName() == null - || metric.getName().getContext() == null - || !metric.getName().getContext().containsKey(TENTATIVE_COUNTER)) { - // Don't double count using the non-tentative version of the metric. - continue; - } - if (PAssert.SUCCESS_COUNTER.equals(metric.getName().getName())) { - successes += ((BigDecimal) metric.getScalar()).intValue(); - } else if (PAssert.FAILURE_COUNTER.equals(metric.getName().getName())) { - failures += ((BigDecimal) metric.getScalar()).intValue(); - } - } + PAssertCounts counts = countPAssertCounters(metrics, step -> true); + int successes = counts.successes; + int failures = counts.failures; if (failures > 0) { LOG.info( @@ -315,6 +488,68 @@ Optional checkForPAssertSuccess(DataflowPipelineJob job) { return Optional.absent(); } + /** Totals of the tentative {@link PAssert} success and failure counters reported by a job. */ + static final class PAssertCounts { + final int successes; + final int failures; + + PAssertCounts(int successes, int failures) { + this.successes = successes; + this.failures = failures; + } + + /** Whether no assertion failed and at least {@code expectedAssertions} succeeded. */ + boolean satisfy(int expectedAssertions) { + return failures == 0 && successes >= expectedAssertions; + } + + @Override + public String toString() { + return successes + " success, " + failures + " failures"; + } + } + + /** + * Sums the {@link PAssert} success and failure counters in {@code metrics}, considering only the + * updates whose Dataflow step (the {@code step} entry of the metric context, or the empty string + * if absent) is accepted by {@code stepFilter}. Only the tentative flavour of each counter is + * counted so that its committed duplicate does not double count. + */ + static PAssertCounts countPAssertCounters(JobMetrics metrics, Predicate stepFilter) { + int successes = 0; + int failures = 0; + List updates = metrics.getMetrics(); + if (updates == null) { + return new PAssertCounts(successes, failures); + } + for (MetricUpdate metric : updates) { + MetricStructuredName name = metric.getName(); + if (name == null) { + continue; + } + boolean isSuccess = PAssert.SUCCESS_COUNTER.equals(name.getName()); + boolean isFailure = PAssert.FAILURE_COUNTER.equals(name.getName()); + if (!isSuccess && !isFailure) { + continue; + } + Map context = name.getContext(); + if (context == null || !context.containsKey(TENTATIVE_COUNTER)) { + continue; + } + String step = context.get("step"); + if (!stepFilter.test(step == null ? "" : step)) { + continue; + } + int value = ((BigDecimal) metric.getScalar()).intValue(); + if (isSuccess) { + successes += value; + } else { + failures += value; + } + } + return new PAssertCounts(successes, failures); + } + @VisibleForTesting @Nullable JobMetrics getJobMetrics(DataflowPipelineJob job) { JobMetrics metrics = null; @@ -336,14 +571,13 @@ public String toString() { * *

Creates an error message representing the concatenation of all error messages seen. */ - private static class ErrorMonitorMessagesHandler implements JobMessagesHandler { + static class ErrorMonitorMessagesHandler implements JobMessagesHandler { private final DataflowPipelineJob job; private final JobMessagesHandler messageHandler; private final StringBuilder errorMessage; private volatile boolean hasSeenError; - private ErrorMonitorMessagesHandler( - DataflowPipelineJob job, JobMessagesHandler messageHandler) { + ErrorMonitorMessagesHandler(DataflowPipelineJob job, JobMessagesHandler messageHandler) { this.job = job; this.messageHandler = messageHandler; this.errorMessage = new StringBuilder(); @@ -374,12 +608,16 @@ String getErrorMessage() { } } - private static class CancelOnError implements Callable { + /** + * Polls a job and cancels it as soon as {@code messageHandler} has seen an error. Used for + * streaming jobs, where a failing step retries forever instead of failing the job. + */ + static class CancelOnError implements Callable { private final DataflowPipelineJob job; private final ErrorMonitorMessagesHandler messageHandler; - public CancelOnError(DataflowPipelineJob job, ErrorMonitorMessagesHandler messageHandler) { + CancelOnError(DataflowPipelineJob job, ErrorMonitorMessagesHandler messageHandler) { this.job = job; this.messageHandler = messageHandler; } diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java index b196aa1714c6..72ca2d7f0a4d 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java @@ -41,7 +41,6 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; -import static org.junit.Assume.assumeFalse; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.eq; @@ -145,8 +144,6 @@ import org.apache.beam.sdk.testing.ExpectedLogs; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.testing.UsesStatefulParDo; -import org.apache.beam.sdk.testing.ValidatesRunner; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.GroupIntoBatches; @@ -154,7 +151,6 @@ import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Redistribute; -import org.apache.beam.sdk.transforms.Reshuffle; import org.apache.beam.sdk.transforms.SerializableFunctions; import org.apache.beam.sdk.transforms.SimpleFunction; import org.apache.beam.sdk.transforms.resourcehints.ResourceHints; @@ -172,17 +168,11 @@ import org.apache.beam.sdk.util.construction.PipelineTranslation; import org.apache.beam.sdk.util.construction.SdkComponents; import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar; -import org.apache.beam.sdk.values.CausedByDrain; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.PValues; import org.apache.beam.sdk.values.TimestampedValue; -import org.apache.beam.sdk.values.TupleTag; -import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.TypeDescriptors; -import org.apache.beam.sdk.values.ValueKind; -import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; @@ -198,10 +188,8 @@ import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.Before; -import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; -import org.junit.experimental.categories.Category; import org.junit.rules.ExpectedException; import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; @@ -232,7 +220,6 @@ public class DataflowRunnerTest implements Serializable { @Rule public transient TemporaryFolder tmpFolder = new TemporaryFolder(); @Rule public transient ExpectedException thrown = ExpectedException.none(); @Rule public transient ExpectedLogs expectedLogs = ExpectedLogs.none(DataflowRunner.class); - @Rule public final transient TestPipeline pipeline = TestPipeline.create(); private transient Dataflow.Projects.Locations.Jobs mockJobs; private transient GcsUtil mockGcsUtil; @@ -2187,7 +2174,11 @@ public void testMergingStatefulRejectedInBatch() throws Exception { verifyMergingStatefulParDoRejected(options); } - private void verifyGroupIntoBatchesOverrideCount( + /** + * Builds a {@link GroupIntoBatches} pipeline on {@code p}, runs it and checks whether the runner + * substituted its override. Shared with {@link DataflowRunnerValidatesRunnerTest}. + */ + static void verifyGroupIntoBatchesOverrideCount( Pipeline p, Boolean withShardedKey, Boolean expectOverridden) { final int batchSize = 2; List> testValues = @@ -2265,7 +2256,8 @@ public CompositeBehavior enterCompositeTransform(Node node) { } } - private void verifyGroupIntoBatchesOverrideBytes( + /** Byte-size counterpart of {@link #verifyGroupIntoBatchesOverrideCount}. */ + static void verifyGroupIntoBatchesOverrideBytes( Pipeline p, Boolean withShardedKey, Boolean expectOverridden) { final long batchSizeBytes = 2; List> testValues = @@ -2339,22 +2331,6 @@ public CompositeBehavior enterCompositeTransform(Node node) { } } - @Test - @Category({ValidatesRunner.class, UsesStatefulParDo.class}) - public void testBatchGroupIntoBatchesOverrideCount() { - // Ignore this test for streaming pipelines. - assumeFalse(pipeline.getOptions().as(StreamingOptions.class).isStreaming()); - verifyGroupIntoBatchesOverrideCount(pipeline, false, true); - } - - @Test - @Category({ValidatesRunner.class, UsesStatefulParDo.class}) - public void testBatchGroupIntoBatchesOverrideBytes() { - // Ignore this test for streaming pipelines. - assumeFalse(pipeline.getOptions().as(StreamingOptions.class).isStreaming()); - verifyGroupIntoBatchesOverrideBytes(pipeline, false, true); - } - @Test public void testBatchGroupIntoBatchesWithShardedKeyOverrideCount() throws IOException { PipelineOptions options = buildPipelineOptions(); @@ -2809,140 +2785,6 @@ public void process() {} p.run(); } - @Test - @Category({ValidatesRunner.class}) - public void testValueKindParameterAndOutputWithKind() { - boolean isRunnerV2 = false; - @Nullable List experiments = - pipeline.getOptions().as(DataflowPipelineOptions.class).getExperiments(); - if (experiments != null - && (experiments.contains("use_unified_worker") || experiments.contains("use_runner_v2"))) { - isRunnerV2 = true; - } - // Skipp runner v2 because its Create uses a splittable DoFn, which contains a shuffle. - // ValueKind is not supported in Dataflow shuffle yet - assumeFalse(isRunnerV2); - - PCollection input = pipeline.apply(Create.of("a", "b", "c", "d")); - TupleTag mainTag = new TupleTag() {}; - TupleTag sideTag = new TupleTag() {}; - - PCollectionTuple tuple = - input.apply( - "SetKind", - ParDo.of( - new DoFn() { - @ProcessElement - public void processElement( - @Element String element, - @Timestamp org.joda.time.Instant timestamp, - BoundedWindow window, - PaneInfo paneInfo, - ProcessContext c, - MultiOutputReceiver outputReceiver) { - switch (element) { - case "a": - c.output(element); // default: INSERT - return; - case "b": - c.outputWindowedValue( - WindowedValues.of( - element, - timestamp, - Collections.singleton(window), - paneInfo, - null, - null, - CausedByDrain.NORMAL, - null, - ValueKind.UPDATE_BEFORE)); - return; - case "c": - outputReceiver - .get(mainTag) - .builder(element) - .setValueKind(ValueKind.UPDATE_AFTER) - .output(); - return; - case "d": - outputReceiver - .get(sideTag) - .builder(element) - .setValueKind(ValueKind.DELETE) - .output(); - } - } - }) - .withOutputTags(mainTag, TupleTagList.of(sideTag))); - - PCollection main = - tuple - .get(mainTag) - .apply( - "ReadKind", - ParDo.of( - new DoFn() { - @ProcessElement - public void processElement( - @Element String element, ProcessContext c, ValueKind kind) { - c.output(element + ":" + kind); - } - })); - - PCollection side = - tuple - .get(sideTag) - .apply( - "ReadKind-SideTag", - ParDo.of( - new DoFn() { - @ProcessElement - public void processElement( - @Element String element, ProcessContext c, ValueKind kind) { - c.output(element + ":" + kind); - } - })); - - PAssert.that(main).containsInAnyOrder("a:INSERT", "b:UPDATE_BEFORE", "c:UPDATE_AFTER"); - PAssert.that(side).containsInAnyOrder("d:DELETE"); - pipeline.run(); - } - - @Test - @Ignore("enable once when element metadata is supported in shuffle") - @Category({ValidatesRunner.class}) - public void testValueKindPreservedAcrossShuffle() { - PCollection> input = pipeline.apply(Create.of(KV.of("key", "value"))); - - PCollection output = - input - .apply( - "SetKind", - ParDo.of( - new DoFn, KV>() { - @ProcessElement - public void processElement( - @Element KV element, - OutputReceiver> out) { - out.builder(element).setValueKind(ValueKind.UPDATE_BEFORE).output(); - } - })) - .apply(Reshuffle.of()) - .apply( - "ReadKind", - ParDo.of( - new DoFn, String>() { - @ProcessElement - public void processElement( - @Element KV element, ProcessContext c, ValueKind kind) { - c.output(element.getValue() + ":" + kind); - } - })); - - PAssert.that(output).containsInAnyOrder("value:UPDATE_BEFORE"); - pipeline.run(); - } - @Test public void testStreamingStateTagEncodingV2PreCompatibility() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerValidatesRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerValidatesRunnerTest.java new file mode 100644 index 000000000000..24ac1c21717f --- /dev/null +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerValidatesRunnerTest.java @@ -0,0 +1,218 @@ +/* + * 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.dataflow; + +import static org.junit.Assume.assumeFalse; + +import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; +import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.testing.BeamParallelJunit4Runner; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.UsesStatefulParDo; +import org.apache.beam.sdk.testing.ValidatesRunner; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.CausedByDrain; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.sdk.values.WindowedValues; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +/** + * {@link ValidatesRunner} tests for {@link DataflowRunner} that execute a pipeline on the service. + * + *

These live apart from {@link DataflowRunnerTest} so that they can run under {@link + * BeamParallelJunit4Runner} (and thus concurrently, and merged into a single Dataflow job by {@link + * TestDataflowRunner}) while the unit tests in {@link DataflowRunnerTest}, some of which rely on + * process-wide state such as log capture, keep running serially under plain JUnit4. + * + *

{@link Serializable} because the anonymous {@link TupleTag}s and {@link DoFn}s below capture + * the enclosing instance. + */ +@RunWith(BeamParallelJunit4Runner.class) +public class DataflowRunnerValidatesRunnerTest implements Serializable { + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @Test + @Category({ValidatesRunner.class, UsesStatefulParDo.class}) + public void testBatchGroupIntoBatchesOverrideCount() { + // Ignore this test for streaming pipelines. + assumeFalse(pipeline.getOptions().as(StreamingOptions.class).isStreaming()); + DataflowRunnerTest.verifyGroupIntoBatchesOverrideCount(pipeline, false, true); + } + + @Test + @Category({ValidatesRunner.class, UsesStatefulParDo.class}) + public void testBatchGroupIntoBatchesOverrideBytes() { + // Ignore this test for streaming pipelines. + assumeFalse(pipeline.getOptions().as(StreamingOptions.class).isStreaming()); + DataflowRunnerTest.verifyGroupIntoBatchesOverrideBytes(pipeline, false, true); + } + + @Test + @Category({ValidatesRunner.class}) + public void testValueKindParameterAndOutputWithKind() { + boolean isRunnerV2 = false; + @Nullable List experiments = + pipeline.getOptions().as(DataflowPipelineOptions.class).getExperiments(); + if (experiments != null + && (experiments.contains("use_unified_worker") || experiments.contains("use_runner_v2"))) { + isRunnerV2 = true; + } + // Skipp runner v2 because its Create uses a splittable DoFn, which contains a shuffle. + // ValueKind is not supported in Dataflow shuffle yet + assumeFalse(isRunnerV2); + + PCollection input = pipeline.apply(Create.of("a", "b", "c", "d")); + TupleTag mainTag = new TupleTag() {}; + TupleTag sideTag = new TupleTag() {}; + + PCollectionTuple tuple = + input.apply( + "SetKind", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element String element, + @Timestamp org.joda.time.Instant timestamp, + BoundedWindow window, + PaneInfo paneInfo, + ProcessContext c, + MultiOutputReceiver outputReceiver) { + switch (element) { + case "a": + c.output(element); // default: INSERT + return; + case "b": + c.outputWindowedValue( + WindowedValues.of( + element, + timestamp, + Collections.singleton(window), + paneInfo, + null, + null, + CausedByDrain.NORMAL, + null, + ValueKind.UPDATE_BEFORE)); + return; + case "c": + outputReceiver + .get(mainTag) + .builder(element) + .setValueKind(ValueKind.UPDATE_AFTER) + .output(); + return; + case "d": + outputReceiver + .get(sideTag) + .builder(element) + .setValueKind(ValueKind.DELETE) + .output(); + } + } + }) + .withOutputTags(mainTag, TupleTagList.of(sideTag))); + + PCollection main = + tuple + .get(mainTag) + .apply( + "ReadKind", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element String element, ProcessContext c, ValueKind kind) { + c.output(element + ":" + kind); + } + })); + + PCollection side = + tuple + .get(sideTag) + .apply( + "ReadKind-SideTag", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element String element, ProcessContext c, ValueKind kind) { + c.output(element + ":" + kind); + } + })); + + PAssert.that(main).containsInAnyOrder("a:INSERT", "b:UPDATE_BEFORE", "c:UPDATE_AFTER"); + PAssert.that(side).containsInAnyOrder("d:DELETE"); + pipeline.run(); + } + + @Test + @Ignore("enable once when element metadata is supported in shuffle") + @Category({ValidatesRunner.class}) + public void testValueKindPreservedAcrossShuffle() { + PCollection> input = pipeline.apply(Create.of(KV.of("key", "value"))); + + PCollection output = + input + .apply( + "SetKind", + ParDo.of( + new DoFn, KV>() { + @ProcessElement + public void processElement( + @Element KV element, + OutputReceiver> out) { + out.builder(element).setValueKind(ValueKind.UPDATE_BEFORE).output(); + } + })) + .apply(Reshuffle.of()) + .apply( + "ReadKind", + ParDo.of( + new DoFn, String>() { + @ProcessElement + public void processElement( + @Element KV element, ProcessContext c, ValueKind kind) { + c.output(element.getValue() + ":" + kind); + } + })); + + PAssert.that(output).containsInAnyOrder("value:UPDATE_BEFORE"); + pipeline.run(); + } +} diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowTestBatchCoordinatorTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowTestBatchCoordinatorTest.java new file mode 100644 index 000000000000..2657a2d787aa --- /dev/null +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowTestBatchCoordinatorTest.java @@ -0,0 +1,1477 @@ +/* + * 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.dataflow; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.isEmptyString; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.startsWith; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.api.services.dataflow.model.ComponentTransform; +import com.google.api.services.dataflow.model.ExecutionStageState; +import com.google.api.services.dataflow.model.ExecutionStageSummary; +import com.google.api.services.dataflow.model.Job; +import com.google.api.services.dataflow.model.JobMetrics; +import com.google.api.services.dataflow.model.MetricStructuredName; +import com.google.api.services.dataflow.model.MetricUpdate; +import com.google.api.services.dataflow.model.PipelineDescription; +import java.io.IOException; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.dataflow.DataflowPipelineTranslator.JobSpecification; +import org.apache.beam.runners.dataflow.DataflowTestBatchCoordinator.BatchingConfig; +import org.apache.beam.runners.dataflow.DataflowTestBatchCoordinator.ScopedDataflowPipelineJob; +import org.apache.beam.runners.dataflow.DataflowTestBatchCoordinator.StageAttribution; +import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions; +import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; +import org.apache.beam.sdk.CompositePipeline; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.PipelineResult.State; +import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; +import org.apache.beam.sdk.extensions.gcp.storage.NoopPathValidator; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.io.GenerateSequence; +import org.apache.beam.sdk.metrics.MetricKey; +import org.apache.beam.sdk.metrics.MetricName; +import org.apache.beam.sdk.metrics.MetricNameFilter; +import org.apache.beam.sdk.metrics.MetricQueryResults; +import org.apache.beam.sdk.metrics.MetricResult; +import org.apache.beam.sdk.metrics.MetricsFilter; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.SerializableMatcher; +import org.apache.beam.sdk.testing.SerializableMatchers; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.util.common.ReflectHelpers; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.SdkComponents; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.hamcrest.Matchers; +import org.joda.time.Duration; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; + +/** + * Tests for {@link DataflowTestBatchCoordinator}, including Dataflow translation of a {@link + * CompositePipeline} (its runner-independent behavior is covered in {@code CompositePipelineTest}). + */ +@RunWith(JUnit4.class) +public class DataflowTestBatchCoordinatorTest { + + private static final long WAIT_SECONDS = 60; + private static final long NEVER_EXPIRING_WINDOW_MS = TimeUnit.MINUTES.toMillis(10); + private static final String STAGE_DONE = "JOB_STATE_DONE"; + private static final String STAGE_FAILED = "JOB_STATE_FAILED"; + private static final String STAGE_PENDING = "JOB_STATE_PENDING"; + + @Mock private DataflowClient mockClient; + private TestDataflowPipelineOptions options; + private ExecutorService pool; + private DataflowTestBatchCoordinator coordinator; + + /** Metrics served by {@link #mockClient}, keyed by job id. Safe to populate from any thread. */ + private final Map metricsByJobId = new ConcurrentHashMap<>(); + + /** The {@code JOB_VIEW_ALL} job served by {@link #mockClient} for any job id. */ + private final AtomicReference executionDetails = new AtomicReference<>(); + + private static class IdentityFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) { + c.output(c.element()); + } + } + + @Before + public void setUp() throws Exception { + MockitoAnnotations.initMocks(this); + options = createTestOptions("TestBatchApp"); + pool = Executors.newCachedThreadPool(); + // Far longer than WAIT_SECONDS: batches in these tests close by reaching maxBatchSize, never + // by the window expiring, so member submissions can't race the timer. Tests that exercise + // window expiry configure their own short window. + configureBatching(2, NEVER_EXPIRING_WINDOW_MS); + when(mockClient.getJobMetrics(anyString())) + .thenAnswer(inv -> metricsByJobId.get(inv.getArgument(0))); + when(mockClient.getJob(anyString(), eq("JOB_VIEW_ALL"))) + .thenAnswer(inv -> executionDetails.get()); + } + + @After + public void tearDown() { + pool.shutdownNow(); + coordinator.shutdown(); + } + + /** + * Gives the current test a private coordinator with the given run-wide configuration. Must be + * called before any {@link Member} is created. + */ + private void configureBatching(int maxBatchSize, long windowMs) { + if (coordinator != null) { + coordinator.shutdown(); + } + coordinator = + new DataflowTestBatchCoordinator( + new BatchingConfig(/* enabled= */ true, maxBatchSize, windowMs)); + } + + private static TestDataflowPipelineOptions createTestOptions(String appName) { + TestDataflowPipelineOptions opts = PipelineOptionsFactory.as(TestDataflowPipelineOptions.class); + opts.setAppName(appName); + opts.setProject("test-project"); + opts.setRegion("us-central1"); + opts.setTempLocation("gs://test/temp/location"); + opts.setTempRoot("gs://test"); + opts.setGcpCredential(new TestCredential()); + opts.setRunner(TestDataflowRunner.class); + opts.setPathValidatorClass(NoopPathValidator.class); + // Pipeline.create registers file systems as a side effect but TestPipeline.fromOptions does + // not, and DataflowRunner.fromOptions needs the gs:// scheme to derive the staging location. + FileSystems.setDefaultPipelineOptions(opts); + return opts; + } + + /** One test pipeline together with the options and runner it is submitted through. */ + private final class Member { + final TestDataflowPipelineOptions opts; + final TestPipeline pipeline; + final TestDataflowRunner runner; + + Member(String appName) { + this.opts = createTestOptions(appName); + this.pipeline = namedTestPipeline(opts, appName); + this.runner = TestDataflowRunner.fromOptionsAndClient(opts, mockClient, coordinator); + } + + /** Applies {@code Create -> ParDo(stepName)} and a PAssert expecting {@code expected}. */ + Member withStep(String stepName, Integer... expected) { + PCollection out = + pipeline + .apply("Create-" + stepName, Create.of(1, 2)) + .apply(stepName, ParDo.of(new IdentityFn())); + PAssert.that(out).containsInAnyOrder(expected); + return this; + } + + Future submit(DataflowRunner delegate) { + return pool.submit(() -> runner.run(pipeline, delegate)); + } + } + + /** + * A {@link TestPipeline} with a unique root name, as {@link TestPipeline}'s JUnit rule produces + * when {@code -DbeamTestPipelineUniqueRootNames=true}. + */ + private static TestPipeline namedTestPipeline(TestDataflowPipelineOptions opts, String testName) { + TestPipeline pipeline = TestPipeline.fromOptions(opts); + pipeline.nameRoot(testName); + assertThat(pipeline.getRootName(), not(isEmptyString())); + return pipeline; + } + + // --------------------------------------------------------------------------------------------- + // Composite pipeline construction + // --------------------------------------------------------------------------------------------- + + @Test + public void testCompositePipelineTranslatesToDataflowWithScopedNames() { + TestPipeline p0 = namedTestPipeline(createTestOptions("App0"), "App0"); + PCollection out0 = + p0.apply("CreateValues", Create.of(1, 2, 3)).apply("MyStep", ParDo.of(new IdentityFn())); + PAssert.that(out0).containsInAnyOrder(1, 2, 3); + + TestPipeline p1 = namedTestPipeline(createTestOptions("App1"), "App1"); + PCollection out1 = + p1.apply("CreateValues", Create.of(4, 5)).apply("MyStep", ParDo.of(new IdentityFn())); + PAssert.that(out1).containsInAnyOrder(4, 5); + + String root0 = p0.getRootName(); + String root1 = p1.getRootName(); + assertThat(root0, not(equalTo(root1))); + // Root names also scope PCollection names, which Dataflow uses for output step names. + assertThat(out0.getName(), startsWith(root0 + "/")); + assertThat(out1.getName(), startsWith(root1 + "/")); + + CompositePipeline composite = new CompositePipeline(options, ImmutableList.of(p0, p1)); + // The coordinator attributes PAsserts per member, never on the composite as a whole. + assertEquals(1, PAssert.countAsserts(p0)); + assertEquals(1, PAssert.countAsserts(p1)); + + DataflowRunner runner = DataflowRunner.fromOptions(options); + SdkComponents portableComponents = + SdkComponents.create(options, Environments.createDockerEnvironment("test-image")); + RunnerApi.Pipeline portableProto = + PipelineTranslation.toProto(composite, portableComponents, false); + + Set uniqueNames = new HashSet<>(); + for (RunnerApi.PTransform transform : + portableProto.getComponents().getTransformsMap().values()) { + uniqueNames.add(transform.getUniqueName()); + } + assertThat(uniqueNames, hasItem(root0 + "/MyStep")); + assertThat(uniqueNames, hasItem(root1 + "/MyStep")); + assertThat(uniqueNames, not(hasItem("MyStep"))); + + runner.replaceV1Transforms(composite); + SdkComponents v1Components = + SdkComponents.create(options, Environments.createDockerEnvironment("test-image")); + RunnerApi.Pipeline v1Proto = PipelineTranslation.toProto(composite, v1Components, true, false); + DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); + JobSpecification jobSpec = + translator.translate(composite, v1Proto, v1Components, runner, Collections.emptyList()); + assertTrue(jobSpec.getStepNames().size() >= 2); + + // Nothing about a member changes by having been part of a composite: it translates on its own + // with exactly the same names, so a standalone fallback run is equivalent to a fresh run. + RunnerApi.Pipeline aloneProto = PipelineTranslation.toProto(p0); + Set aloneUniqueNames = new HashSet<>(); + for (RunnerApi.PTransform transform : aloneProto.getComponents().getTransformsMap().values()) { + aloneUniqueNames.add(transform.getUniqueName()); + } + assertThat(aloneUniqueNames, hasItem(root0 + "/MyStep")); + assertEquals(root0, p0.getRootName()); + } + + @Test + public void testEligibilityChecks() { + TestPipeline boundedPipeline = namedTestPipeline(options, "Bounded"); + boundedPipeline.apply(Create.of(1, 2, 3)); + assertTrue(coordinator.isEligibleForBatching(boundedPipeline, options)); + + // Streaming pipelines bypass batching. + TestDataflowPipelineOptions streamingOpts = createTestOptions("StreamingApp"); + streamingOpts.setStreaming(true); + assertFalse(coordinator.isEligibleForBatching(boundedPipeline, streamingOpts)); + + // Unbounded pipelines bypass batching. + TestPipeline unboundedPipeline = namedTestPipeline(options, "Unbounded"); + unboundedPipeline.apply(GenerateSequence.from(0)); + assertFalse(coordinator.isEligibleForBatching(unboundedPipeline, options)); + + // Only TestPipelines with a unique root name can be told apart inside a merged job. + Pipeline plainPipeline = Pipeline.create(options); + plainPipeline.apply(Create.of(1, 2, 3)); + assertFalse(coordinator.isEligibleForBatching(plainPipeline, options)); + TestPipeline unnamedPipeline = TestPipeline.fromOptions(options); + unnamedPipeline.apply(Create.of(1, 2, 3)); + assertThat(unnamedPipeline.getRootName(), isEmptyString()); + assertFalse(coordinator.isEligibleForBatching(unnamedPipeline, options)); + + // Batching is a run-wide switch (coordinator config), not a per-pipeline option. + DataflowTestBatchCoordinator disabled = + new DataflowTestBatchCoordinator(new BatchingConfig(false, 2, 5000L)); + try { + assertFalse(disabled.isEligibleForBatching(boundedPipeline, options)); + } finally { + disabled.shutdown(); + } + DataflowTestBatchCoordinator singletonBatches = + new DataflowTestBatchCoordinator(new BatchingConfig(true, 1, 5000L)); + try { + assertFalse(singletonBatches.isEligibleForBatching(boundedPipeline, options)); + } finally { + singletonBatches.shutdown(); + } + + // A non-blocking run has no verdict to wait for. + TestDataflowPipelineOptions nonBlocking = createTestOptions("NonBlocking"); + nonBlocking.setBlockOnRun(false); + assertFalse(coordinator.isEligibleForBatching(boundedPipeline, nonBlocking)); + + // Custom result matchers would be evaluated against the merged job, not the member's own. + SerializableMatcher customMatcher = + SerializableMatchers.fromSupplier(() -> Matchers.notNullValue(PipelineResult.class)); + TestDataflowPipelineOptions successMatcher = createTestOptions("SuccessMatcher"); + successMatcher.setOnSuccessMatcher(customMatcher); + assertFalse(coordinator.isEligibleForBatching(boundedPipeline, successMatcher)); + TestDataflowPipelineOptions createMatcher = createTestOptions("CreateMatcher"); + createMatcher.setOnCreateMatcher(customMatcher); + assertFalse(coordinator.isEligibleForBatching(boundedPipeline, createMatcher)); + + // Whatever TestPipeline decides requires standalone execution (expected failures, serial + // tests, runtime value providers, ...) is honoured; newProvider is the simplest trigger. + TestPipeline standalone = namedTestPipeline(options, "Standalone"); + standalone.apply(Create.of(1, 2, 3)); + standalone.newProvider("runtime-value"); + assertTrue(standalone.isStandaloneExecutionRequired()); + assertFalse(coordinator.isEligibleForBatching(standalone, options)); + + // A pipeline with no primitive transform has nothing to run or verify. + assertFalse(coordinator.isEligibleForBatching(namedTestPipeline(options, "Empty"), options)); + + // An already-merged pipeline is never merged again. + assertFalse( + coordinator.isEligibleForBatching( + new CompositePipeline(options, ImmutableList.of(boundedPipeline)), options)); + } + + // --------------------------------------------------------------------------------------------- + // Options compatibility + // --------------------------------------------------------------------------------------------- + + /** An options interface the coordinator knows nothing about. */ + public interface CustomTestOptions extends PipelineOptions { + @Nullable String getFlavor(); + + void setFlavor(String value); + + @Nullable Object getOpaque(); + + void setOpaque(Object value); + } + + @Test + public void testCompatibilityKeyIgnoresPerTestIdentityOptions() { + TestDataflowPipelineOptions a = createTestOptions("AppA"); + TestDataflowPipelineOptions b = createTestOptions("AppB"); + b.setJobName("some-other-job-name"); + b.setTempLocation("gs://test/other/temp"); + b.setGcpTempLocation("gs://test/other/gcptemp"); + b.setOptionsId(a.getOptionsId() + 1); + b.setTestTimeoutSeconds(42L); + assertNotNull(DataflowTestBatchCoordinator.compatibilityKey(a)); + assertEquals( + DataflowTestBatchCoordinator.compatibilityKey(a), + DataflowTestBatchCoordinator.compatibilityKey(b)); + } + + @Test + public void testCompatibilityKeyDiffersForAnyJobAffectingOption() { + String base = DataflowTestBatchCoordinator.compatibilityKey(createTestOptions("App")); + + TestDataflowPipelineOptions machineType = createTestOptions("App"); + machineType.as(DataflowPipelineWorkerPoolOptions.class).setWorkerMachineType("n2-standard-8"); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(machineType)); + + TestDataflowPipelineOptions experiments = createTestOptions("App"); + experiments.setExperiments(ImmutableList.of("use_runner_v2")); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(experiments)); + + TestDataflowPipelineOptions debug = createTestOptions("App"); + debug.as(DataflowPipelineDebugOptions.class).setNumberOfWorkerHarnessThreads(3); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(debug)); + + TestDataflowPipelineOptions region = createTestOptions("App"); + region.setRegion("europe-west1"); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(region)); + + TestDataflowPipelineOptions streaming = createTestOptions("App"); + streaming.setStreaming(true); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(streaming)); + + // Options set through an interface the coordinator has never heard of still count. + TestDataflowPipelineOptions vanilla = createTestOptions("App"); + vanilla.as(CustomTestOptions.class).setFlavor("vanilla"); + TestDataflowPipelineOptions vanillaAgain = createTestOptions("App"); + vanillaAgain.as(CustomTestOptions.class).setFlavor("vanilla"); + TestDataflowPipelineOptions chocolate = createTestOptions("App"); + chocolate.as(CustomTestOptions.class).setFlavor("chocolate"); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(vanilla)); + assertEquals( + DataflowTestBatchCoordinator.compatibilityKey(vanilla), + DataflowTestBatchCoordinator.compatibilityKey(vanillaAgain)); + assertNotEquals( + DataflowTestBatchCoordinator.compatibilityKey(vanilla), + DataflowTestBatchCoordinator.compatibilityKey(chocolate)); + } + + @Test + public void testCompatibilityKeyIsNullForUnserializableOptions() { + TestDataflowPipelineOptions opts = createTestOptions("App"); + opts.as(CustomTestOptions.class).setOpaque(new Object()); + assertNull(DataflowTestBatchCoordinator.compatibilityKey(opts)); + } + + @Test + public void testCompatibilityKeyIgnoresDefaultsBoundByGetters() { + TestDataflowPipelineOptions untouched = createTestOptions("App"); + String base = DataflowTestBatchCoordinator.compatibilityKey(untouched); + + // What View.asList() and Reshuffle do during expansion: read an option nobody has set. The + // getter binds the default, which the serialized form then includes; the key must not change. + TestDataflowPipelineOptions read = createTestOptions("App"); + assertNull(read.as(StreamingOptions.class).getUpdateCompatibilityVersion()); + assertEquals(0, read.as(DataflowPipelineWorkerPoolOptions.class).getNumWorkers()); + assertEquals(base, DataflowTestBatchCoordinator.compatibilityKey(read)); + + // Explicitly setting a value, even one equal to the default, still counts: it is what the test + // asked for and the merged job must honour it. + TestDataflowPipelineOptions set = createTestOptions("App"); + set.as(DataflowPipelineWorkerPoolOptions.class).setNumWorkers(0); + assertNotEquals(base, DataflowTestBatchCoordinator.compatibilityKey(set)); + } + + @Test + public void testCompatibilityKeyHandlesStructuredValuesAcrossTestPipelineRoundTrip() + throws Exception { + // An explicitly set map: its content is part of the key, its insertion order is not. + TestDataflowPipelineOptions ab = createTestOptions("App"); + Map labelsAb = new LinkedHashMap<>(); + labelsAb.put("team", "beam"); + labelsAb.put("tier", "test"); + ab.setLabels(labelsAb); + TestDataflowPipelineOptions ba = createTestOptions("App"); + Map labelsBa = new LinkedHashMap<>(); + labelsBa.put("tier", "test"); + labelsBa.put("team", "beam"); + ba.setLabels(labelsBa); + TestDataflowPipelineOptions other = createTestOptions("App"); + other.setLabels(ImmutableMap.of("team", "beam", "tier", "prod")); + String keyAb = DataflowTestBatchCoordinator.compatibilityKey(ab); + assertEquals(keyAb, DataflowTestBatchCoordinator.compatibilityKey(ba)); + assertNotEquals(keyAb, DataflowTestBatchCoordinator.compatibilityKey(other)); + + // Lists are compared in order: this errs on the side of splitting a batch. + TestDataflowPipelineOptions xy = createTestOptions("App"); + xy.setExperiments(ImmutableList.of("x", "y")); + TestDataflowPipelineOptions yx = createTestOptions("App"); + yx.setExperiments(ImmutableList.of("y", "x")); + assertNotEquals( + DataflowTestBatchCoordinator.compatibilityKey(xy), + DataflowTestBatchCoordinator.compatibilityKey(yx)); + + // A structured default merely read by a getter stays out of the key, and an explicitly set + // structured value stays in it, also on the JSON copy TestPipeline.run hands the runner. + TestDataflowPipelineOptions untouched = createTestOptions("App"); + String base = DataflowTestBatchCoordinator.compatibilityKey(untouched); + TestDataflowPipelineOptions read = createTestOptions("App"); + assertNotNull(read.as(GcpOptions.class).getGcpOauthScopes()); + assertEquals(base, DataflowTestBatchCoordinator.compatibilityKey(read)); + assertEquals(base, DataflowTestBatchCoordinator.compatibilityKey(testPipelineCopy(read))); + assertEquals(keyAb, DataflowTestBatchCoordinator.compatibilityKey(testPipelineCopy(ab))); + assertEquals(keyAb, DataflowTestBatchCoordinator.compatibilityKey(testPipelineCopy(ba))); + assertNotEquals(keyAb, DataflowTestBatchCoordinator.compatibilityKey(testPipelineCopy(other))); + } + + /** The JSON round trip {@code TestPipeline.run(PipelineOptions)} performs before running. */ + private static TestDataflowPipelineOptions testPipelineCopy(PipelineOptions options) + throws IOException { + ObjectMapper mapper = + new ObjectMapper() + .registerModules(ObjectMapper.findModules(ReflectHelpers.findClassLoader())); + byte[] json = + mapper + .writer() + .withAttribute(PipelineOptionsFactory.SERIALIZE_DEFAULTS_ATTRIBUTE, true) + .writeValueAsBytes(options); + return mapper.readValue(json, PipelineOptions.class).as(TestDataflowPipelineOptions.class); + } + + @Test + public void testMembersWithIncompatibleOptionsAreNeverMerged() throws Exception { + // Each member ends up alone in its group; keep the wait for a partner short. + configureBatching(2, 300L); + Member small = new Member("SmallWorker").withStep("StepA", 1, 2); + Member big = new Member("BigWorker").withStep("StepB", 1, 2); + big.opts.as(DataflowPipelineWorkerPoolOptions.class).setWorkerMachineType("n2-standard-8"); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + int callNum = submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat( + "pipelines with different worker options must not share a job", + submitted, + not(instanceOf(CompositePipeline.class))); + String jobId = "standalone-" + callNum; + metricsByJobId.put( + jobId, + new JobMetrics() + .setMetrics( + ImmutableList.of( + createTentativePAssertUpdate("s1", PAssert.SUCCESS_COUNTER, 1)))); + return newStandaloneJob(jobId, State.DONE); + }); + + Future f0 = small.submit(delegate); + Future f1 = big.submit(delegate); + assertEquals(State.DONE, f0.get(WAIT_SECONDS, TimeUnit.SECONDS).getState()); + assertEquals(State.DONE, f1.get(WAIT_SECONDS, TimeUnit.SECONDS).getState()); + assertEquals(2, submittedJobs.get()); + } + + // --------------------------------------------------------------------------------------------- + // Merged job succeeds + // --------------------------------------------------------------------------------------------- + + @Test + public void testConcurrentPipelinesMergedIntoSingleBatchJobWithScopedMetrics() throws Exception { + Member m0 = new Member("App0").withStep("StepA", 1, 2); + Member m1 = new Member("App1").withStep("StepB", 1, 2); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(invocation.getArgument(0)); + List updates = new ArrayList<>(); + updates.addAll( + createCounterUpdates( + transformIdOf(proto, "StepA"), "CustomNs", "myCounter", 10L)); + updates.add(passertUpdate(proto, "StepA", PAssert.SUCCESS_COUNTER)); + updates.addAll( + createCounterUpdates( + transformIdOf(proto, "StepB"), "CustomNs", "myCounter", 20L)); + updates.add(passertUpdate(proto, "StepB", PAssert.SUCCESS_COUNTER)); + metricsByJobId.put("batch-job-1", new JobMetrics().setMetrics(updates)); + return newBatchJob("batch-job-1", State.DONE, proto); + }); + // The spy's DataflowMetrics was constructed around the un-spied job, so metric queries reach + // the real getState(), which polls the client. + when(mockClient.getJob(anyString())).thenReturn(new Job().setCurrentState(STAGE_DONE)); + + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + DataflowPipelineJob result0 = f0.get(WAIT_SECONDS, TimeUnit.SECONDS); + DataflowPipelineJob result1 = f1.get(WAIT_SECONDS, TimeUnit.SECONDS); + + assertEquals(1, submittedJobs.get()); + assertThat(result0, instanceOf(ScopedDataflowPipelineJob.class)); + assertThat(result1, instanceOf(ScopedDataflowPipelineJob.class)); + assertEquals(State.DONE, result0.getState()); + assertEquals(State.DONE, result1.waitUntilFinish()); + + assertSingleCounter(result0, "StepA", 10L); + assertSingleCounter(result1, "StepB", 20L); + // Members are not mutated by the merged run: their root names (and thus graphs) are unchanged. + assertThat(m0.pipeline.getRootName(), startsWith("t")); + assertThat(m0.pipeline.getRootName(), not(equalTo(m1.pipeline.getRootName()))); + } + + @Test + public void testPartialFailureInBatchOnlyRequiresTheFailingTestToBeExecutedAgain() + throws Exception { + configureBatching(3, NEVER_EXPIRING_WINDOW_MS); + Member good0 = new Member("GoodTest0").withStep("GoodStep0", 1, 2); + Member failing1 = new Member("FailingTest1").withStep("FailingStep1", 99); + Member good2 = new Member("GoodTest2").withStep("GoodStep2", 1, 2); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + metricsByJobId.put( + "batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "GoodStep0", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "FailingStep1", PAssert.FAILURE_COUNTER), + passertUpdate(proto, "GoodStep2", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("batch-job", State.DONE, proto); + }); + + Future f0 = good0.submit(delegate); + Future f1 = failing1.submit(delegate); + Future f2 = good2.submit(delegate); + + // GoodTest0 and GoodTest2 are credited directly from the batch job. + assertThat(f0.get(WAIT_SECONDS, TimeUnit.SECONDS), instanceOf(ScopedDataflowPipelineJob.class)); + assertThat(f2.get(WAIT_SECONDS, TimeUnit.SECONDS), instanceOf(ScopedDataflowPipelineJob.class)); + + // FailingTest1 is not run again here: its test gets a request to be executed from scratch. + String reason = rerunReasonOf(f1); + assertThat(reason, containsString("FailingTest1")); + assertThat(reason, containsString("batch-job")); + assertThat(reason, containsString("job DONE, PAsserts unverified")); + + // The merged job is the only one submitted through the coordinator. + assertEquals(1, submittedJobs.get()); + } + + // --------------------------------------------------------------------------------------------- + // Merged job fails + // --------------------------------------------------------------------------------------------- + + @Test + public void testFailedBatchJobWithoutStageStatesRequiresEveryMemberToBeExecutedAgain() + throws Exception { + Member good0 = new Member("GoodTest0").withStep("GoodStep0", 1, 2); + Member failing1 = new Member("FailingTest1").withStep("FailingStep1", 99); + + // executionDetails stays null: the service returned no per-stage information. + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + return newBatchJob( + "failed-batch-job", State.FAILED, PipelineTranslation.toProto(submitted)); + }); + + Future f0 = good0.submit(delegate); + Future f1 = failing1.submit(delegate); + + // Without stage states nobody can be credited, not even the healthy test. + for (Future f : Arrays.asList(f0, f1)) { + String reason = rerunReasonOf(f); + assertThat(reason, containsString("failed-batch-job")); + assertThat(reason, containsString("stage states unavailable")); + } + assertEquals(1, submittedJobs.get()); + } + + @Test + public void testFailedBatchJobCreditsMembersWhoseStagesAllCompleted() throws Exception { + configureBatching(3, NEVER_EXPIRING_WINDOW_MS); + Member good0 = new Member("GoodTest0").withStep("GoodStep0", 1, 2); + Member failing1 = new Member("FailingTest1").withStep("FailingStep1", 99); + Member pending2 = new Member("PendingTest2").withStep("PendingStep2", 1, 2); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + String s0 = scopeOf(proto, "GoodStep0"); + String s1 = scopeOf(proto, "FailingStep1"); + String s2 = scopeOf(proto, "PendingStep2"); + // The service reports per-stage states; only the failing member's stage failed. + // PendingTest2 has a success counter but one of its stages never reached DONE, so it + // must not be credited either. + executionDetails.set( + executionDetails( + stage("F1", STAGE_DONE, s0 + "/Create-GoodStep0/Read", s0 + "/GoodStep0"), + stage("F2", STAGE_FAILED, s1 + "/FailingStep1"), + stage("F3", STAGE_DONE, s2 + "/PendingStep2"), + stage("F4", STAGE_PENDING, s2 + "/Create-PendingStep2/Read"))); + metricsByJobId.put( + "failed-batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "GoodStep0", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "PendingStep2", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("failed-batch-job", State.FAILED, proto); + }); + + Future f0 = good0.submit(delegate); + Future f1 = failing1.submit(delegate); + Future f2 = pending2.submit(delegate); + + DataflowPipelineJob result0 = f0.get(WAIT_SECONDS, TimeUnit.SECONDS); + assertThat(result0, instanceOf(ScopedDataflowPipelineJob.class)); + assertEquals(State.DONE, result0.getState()); + assertEquals(State.DONE, result0.waitUntilFinish()); + + // The others must be executed again from scratch; the reason says why each was not credited. + String reason1 = rerunReasonOf(f1); + assertThat(reason1, containsString("FailingTest1")); + assertThat(reason1, containsString("failed-batch-job")); + assertThat(reason1, containsString("PAsserts unverified")); + String reason2 = rerunReasonOf(f2); + assertThat(reason2, containsString("PendingTest2")); + assertThat(reason2, containsString("PAsserts verified")); + + // Only the merged job was ever submitted through the coordinator. + assertEquals(1, submittedJobs.get()); + } + + @Test + public void testFailedBatchJobWithUnattributableFailedStageCreditsNobody() throws Exception { + Member m0 = new Member("Test0").withStep("Step0", 1, 2); + Member m1 = new Member("Test1").withStep("Step1", 1, 2); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + // Both members' own stages are DONE and both have success counters, but a stage + // that names no user steps failed. It cannot be attributed, so it counts against + // everyone. + executionDetails.set( + executionDetails( + stage("F1", STAGE_DONE, scopeOf(proto, "Step0") + "/Step0"), + stage("F2", STAGE_DONE, scopeOf(proto, "Step1") + "/Step1"), + stage("F99", STAGE_FAILED))); + metricsByJobId.put( + "failed-batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "Step0", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "Step1", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("failed-batch-job", State.FAILED, proto); + }); + + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + assertThat(rerunReasonOf(f0), containsString("job FAILED")); + assertThat(rerunReasonOf(f1), containsString("job FAILED")); + assertEquals(1, submittedJobs.get()); + } + + @Test + public void testStageAttributionMatching() { + List scopes = ImmutableList.of("t0", "t1", "t10"); + StageAttribution attribution = + StageAttribution.fromJob( + executionDetails( + stage("F1", STAGE_DONE, "t0/A"), + stage("F2", STAGE_FAILED, "t1/B"), + // Runner V2 reports mangled names using '-' instead of '/'. + stage("F3", STAGE_DONE, "t0-A-out0"), + // "t10/..." must not be mistaken for scope "t1". + stage("F4", STAGE_FAILED, "t10/C"))); + assertNotNull(attribution); + assertTrue(attribution.statusForScope("t0", scopes).isComplete()); + assertFalse(attribution.statusForScope("t1", scopes).isComplete()); + assertFalse(attribution.statusForScope("t10", scopes).isComplete()); + assertThat(attribution.statusForScope("t0", scopes).toString(), is("2 stage(s) done")); + assertThat(attribution.statusForScope("t1", scopes).toString(), containsString("F2")); + assertThat(attribution.statusForScope("t1", scopes).toString(), not(containsString("F4"))); + + // A scope that owns no stages at all is not complete. + assertFalse(attribution.statusForScope("t7", scopes).isComplete()); + + // A failed stage whose names fall under no known scope counts against every scope. + StageAttribution withUnknown = + StageAttribution.fromJob( + executionDetails( + stage("F1", STAGE_DONE, "t0/A"), stage("F5", STAGE_FAILED, "Mystery/X"))); + assertNotNull(withUnknown); + assertFalse(withUnknown.statusForScope("t0", scopes).isComplete()); + + // A stage with no state entry at all is treated as not done. + StageAttribution missingState = + new StageAttributionBuilder() + .describe("F1", "t0/A") + .describe("F2", "t0/B") + .state("F1", STAGE_DONE) + .build(); + assertNotNull(missingState); + assertFalse(missingState.statusForScope("t0", scopes).isComplete()); + assertThat(missingState.statusForScope("t0", scopes).toString(), containsString("F2=UNKNOWN")); + } + + @Test + public void testStageAttributionRequiresBothStatesAndDescriptions() { + assertNull(StageAttribution.fromJob(null)); + assertNull(StageAttribution.fromJob(new Job())); + assertNull( + StageAttribution.fromJob( + new Job() + .setStageStates( + ImmutableList.of( + new ExecutionStageState() + .setExecutionStageName("F1") + .setExecutionStageState(STAGE_DONE))))); + assertNull( + StageAttribution.fromJob( + new Job() + .setPipelineDescription( + new PipelineDescription() + .setExecutionPipelineStage( + ImmutableList.of(new ExecutionStageSummary().setName("F1")))))); + assertNotNull(StageAttribution.fromJob(executionDetails(stage("F1", STAGE_DONE, "t0/A")))); + } + + // --------------------------------------------------------------------------------------------- + // Submission failures, interruption and unexpected errors + // --------------------------------------------------------------------------------------------- + + @Test + public void testFailedSubmissionRequiresEveryMemberToBeExecutedAgain() throws Exception { + Member m0 = new Member("App0").withStep("Step0", 1, 2); + Member m1 = new Member("App1").withStep("Step1", 1, 2); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + throw new RuntimeException("Simulated batch submission failure"); + }); + + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + + for (Future f : Arrays.asList(f0, f1)) { + String reason = rerunReasonOf(f); + assertThat(reason, containsString("(never submitted)")); + assertThat(reason, containsString("Simulated batch submission failure")); + } + // The member pipelines themselves are never submitted by the coordinator. + assertEquals(1, submittedJobs.get()); + } + + @Test + public void testInterruptedMemberDoesNotAffectOtherMembers() throws Exception { + Member m0 = new Member("App0").withStep("Step0", 1, 2); + Member m1 = new Member("App1").withStep("Step1", 1, 2); + + CountDownLatch compositeEntered = new CountDownLatch(1); + CountDownLatch releaseComposite = new CountDownLatch(1); + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + compositeEntered.countDown(); + assertTrue(releaseComposite.await(WAIT_SECONDS, TimeUnit.SECONDS)); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + metricsByJobId.put( + "batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "Step0", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "Step1", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("batch-job", State.DONE, proto); + }); + + AtomicReference failure0 = new AtomicReference<>(); + Future f0 = + pool.submit( + () -> { + try { + return m0.runner.run(m0.pipeline, delegate); + } catch (RuntimeException e) { + failure0.set(e); + throw e; + } + }); + Future f1 = m1.submit(delegate); + + assertTrue(compositeEntered.await(WAIT_SECONDS, TimeUnit.SECONDS)); + // Member 0's test is interrupted (e.g. by a JUnit timeout) while the merged job is running. + f0.cancel(true); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(WAIT_SECONDS); + while (failure0.get() == null && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertNotNull("member 0 should have observed the interrupt", failure0.get()); + assertThat(failure0.get().getCause(), instanceOf(InterruptedException.class)); + + releaseComposite.countDown(); + DataflowPipelineJob result1 = f1.get(WAIT_SECONDS, TimeUnit.SECONDS); + assertThat(result1, instanceOf(ScopedDataflowPipelineJob.class)); + assertEquals(1, submittedJobs.get()); + } + + @Test + public void testUnexpectedCoordinatorFailureFailsAllMembersPromptly() throws Exception { + Member m0 = new Member("App0").withStep("Step0", 1, 2); + Member m1 = new Member("App1").withStep("Step1", 1, 2); + + // Not an IOException, so it escapes TestDataflowRunner.getJobMetrics and surfaces as an + // unexpected failure inside the batch runner after the merged job has completed. + when(mockClient.getJobMetrics(anyString())) + .thenThrow(new IllegalStateException("metrics backend exploded")); + + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + return newBatchJob("batch-job", State.DONE, PipelineTranslation.toProto(submitted)); + }); + + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + for (Future f : Arrays.asList(f0, f1)) { + Throwable cause = causeOf(f); + assertThat(cause, instanceOf(IllegalStateException.class)); + assertThat(cause.getMessage(), containsString("metrics backend exploded")); + } + // Nobody fell back to a standalone run: the failure was in the test harness, not the job. + assertEquals(1, submittedJobs.get()); + } + + // --------------------------------------------------------------------------------------------- + // Scheduling and concurrency + // --------------------------------------------------------------------------------------------- + + @Test + public void testStragglerIsNotBlockedByInFlightBatch() throws Exception { + // A and B are submitted back-to-back and close a batch well within the window; the straggler + // then waits out one full window alone, so keep it short. + configureBatching(2, 1000L); + Member a = new Member("A").withStep("StepA", 1, 2); + Member b = new Member("B").withStep("StepB", 1, 2); + Member straggler = new Member("Straggler").withStep("StepC", 1, 2); + + CountDownLatch compositeEntered = new CountDownLatch(1); + CountDownLatch releaseComposite = new CountDownLatch(1); + AtomicBoolean stragglerRanWhileBatchInFlight = new AtomicBoolean(false); + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + int callNum = submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + if (submitted instanceof CompositePipeline) { + compositeEntered.countDown(); + assertTrue(releaseComposite.await(WAIT_SECONDS, TimeUnit.SECONDS)); + metricsByJobId.put( + "batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "StepA", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "StepB", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("batch-job", State.DONE, proto); + } + assertTrue(hasTransform(proto, "StepC")); + stragglerRanWhileBatchInFlight.set(releaseComposite.getCount() == 1); + String jobId = "standalone-" + callNum; + metricsByJobId.put( + jobId, + new JobMetrics() + .setMetrics( + ImmutableList.of( + createTentativePAssertUpdate("s1", PAssert.SUCCESS_COUNTER, 1)))); + return newStandaloneJob(jobId, State.DONE); + }); + + Future fa = a.submit(delegate); + Future fb = b.submit(delegate); + assertTrue(compositeEntered.await(WAIT_SECONDS, TimeUnit.SECONDS)); + + // While the {A, B} job is in flight, a third test arrives. Its batch window closes with no + // partner, so it must run standalone right away rather than wait for the in-flight job. + Future fc = straggler.submit(delegate); + DataflowPipelineJob resultC = fc.get(WAIT_SECONDS, TimeUnit.SECONDS); + assertThat(resultC, not(instanceOf(ScopedDataflowPipelineJob.class))); + assertTrue(stragglerRanWhileBatchInFlight.get()); + + releaseComposite.countDown(); + assertThat(fa.get(WAIT_SECONDS, TimeUnit.SECONDS), instanceOf(ScopedDataflowPipelineJob.class)); + assertThat(fb.get(WAIT_SECONDS, TimeUnit.SECONDS), instanceOf(ScopedDataflowPipelineJob.class)); + assertEquals(2, submittedJobs.get()); + } + + @Test + public void testBatchClosesAtMaxSizeWithoutWaitingForWindow() throws Exception { + // The window is far longer than this test is willing to wait, so the merged job can only be + // submitted in time if the batch closes as soon as it reaches maxBatchSize. + long windowMs = NEVER_EXPIRING_WINDOW_MS; + configureBatching(2, windowMs); + Member a = new Member("A").withStep("StepA", 1, 2); + Member b = new Member("B").withStep("StepB", 1, 2); + + CountDownLatch compositeEntered = new CountDownLatch(1); + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + compositeEntered.countDown(); + metricsByJobId.put( + "batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "StepA", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "StepB", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("batch-job", State.DONE, proto); + }); + + long startNanos = System.nanoTime(); + Future fa = a.submit(delegate); + Future fb = b.submit(delegate); + assertTrue(compositeEntered.await(WAIT_SECONDS, TimeUnit.SECONDS)); + long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos); + assertTrue( + "Batch of maxBatchSize members waited " + elapsedMs + "ms, i.e. for the window to expire", + elapsedMs < windowMs / 2); + + assertThat(fa.get(WAIT_SECONDS, TimeUnit.SECONDS), instanceOf(ScopedDataflowPipelineJob.class)); + assertThat(fb.get(WAIT_SECONDS, TimeUnit.SECONDS), instanceOf(ScopedDataflowPipelineJob.class)); + assertEquals(1, submittedJobs.get()); + } + + // --------------------------------------------------------------------------------------------- + // Merged job naming + // --------------------------------------------------------------------------------------------- + + private static final String DATAFLOW_JOB_NAME_REGEX = "[a-z]([-a-z0-9]*[a-z0-9])?"; + + @Test + public void testBatchJobNameIsValidAndKeepsUniqueSuffix() { + // 2026-10-02T16:00:00Z + long now = 1790956800000L; + String name = + DataflowTestBatchCoordinator.batchJobName(7, "ParDoTest$Basic-tests", now, 0xdeadbeef); + assertEquals("batch-7-pardotest0basic0tests-1002160000-deadbeef", name); + assertTrue(name, name.matches(DATAFLOW_JOB_NAME_REGEX)); + + // A long app name (typical for ValidatesRunner suites) is truncated, but only the hint is: + // the timestamp and random suffix that make the name unique always survive. + StringBuilder longAppName = new StringBuilder(); + for (int i = 0; i < 20; i++) { + longAppName.append("SomeVeryLongTestClassName"); + } + String truncated = + DataflowTestBatchCoordinator.batchJobName(7, longAppName.toString(), now, 0xdeadbeef); + assertEquals(DataflowTestBatchCoordinator.MAX_JOB_NAME_LENGTH, truncated.length()); + assertTrue(truncated, truncated.startsWith("batch-7-someverylongtestclassname")); + assertTrue(truncated, truncated.endsWith("-1002160000-deadbeef")); + assertTrue(truncated, truncated.matches(DATAFLOW_JOB_NAME_REGEX)); + + // Same batch number (as in two JVMs running the same suite concurrently) and same app name + // still yield distinct names thanks to the random part. + assertNotEquals( + DataflowTestBatchCoordinator.batchJobName(1, longAppName.toString(), now, 1), + DataflowTestBatchCoordinator.batchJobName(1, longAppName.toString(), now, 2)); + + // Degenerate hints never produce an invalid name. + assertTrue( + DataflowTestBatchCoordinator.batchJobName(1, null, now, 1) + .matches(DATAFLOW_JOB_NAME_REGEX)); + assertTrue( + DataflowTestBatchCoordinator.batchJobName(1, "", now, 1).matches(DATAFLOW_JOB_NAME_REGEX)); + assertTrue( + DataflowTestBatchCoordinator.batchJobName(1, "---", now, 1) + .matches(DATAFLOW_JOB_NAME_REGEX)); + } + + @Test + public void testMergedJobIsSubmittedUnderBatchNameAndLeaderNameIsRestored() throws Exception { + Member m0 = new Member("App0").withStep("Step0", 1, 2); + Member m1 = new Member("App1").withStep("Step1", 1, 2); + String leaderJobName = m0.opts.getJobName(); + String otherJobName = m1.opts.getJobName(); + + AtomicReference submittedName = new AtomicReference<>(); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + submittedName.set(submitted.getOptions().getJobName()); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(submitted); + metricsByJobId.put( + "batch-job", + new JobMetrics() + .setMetrics( + ImmutableList.of( + passertUpdate(proto, "Step0", PAssert.SUCCESS_COUNTER), + passertUpdate(proto, "Step1", PAssert.SUCCESS_COUNTER)))); + return newBatchJob("batch-job", State.DONE, proto); + }); + + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + f0.get(WAIT_SECONDS, TimeUnit.SECONDS); + f1.get(WAIT_SECONDS, TimeUnit.SECONDS); + + String name = submittedName.get(); + assertNotNull(name); + assertTrue(name, name.startsWith("batch-")); + assertTrue(name, name.matches(DATAFLOW_JOB_NAME_REGEX)); + assertTrue(name, name.length() <= DataflowTestBatchCoordinator.MAX_JOB_NAME_LENGTH); + assertNotEquals(leaderJobName, name); + // The leader's own job name is intact afterwards; the batch name was only ever borrowed. + assertEquals(leaderJobName, m0.opts.getJobName()); + assertEquals(otherJobName, m1.opts.getJobName()); + } + + // --------------------------------------------------------------------------------------------- + // No-hang guarantee + // --------------------------------------------------------------------------------------------- + + @Test + public void testMergedJobExceedingShortestMemberTimeoutIsCancelledAndMembersRerun() + throws Exception { + Member m0 = new Member("App0").withStep("Step0", 1, 2); + Member m1 = new Member("App1").withStep("Step1", 1, 2); + // Different per-test timeouts do not prevent merging; the merged job gets the shortest one. + m0.opts.setTestTimeoutSeconds(7L); + m1.opts.setTestTimeoutSeconds(3L); + + AtomicReference batchJob = new AtomicReference<>(); + AtomicReference waitedFor = new AtomicReference<>(); + AtomicInteger submittedJobs = new AtomicInteger(0); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + submittedJobs.incrementAndGet(); + Pipeline submitted = invocation.getArgument(0); + assertThat(submitted, instanceOf(CompositePipeline.class)); + DataflowPipelineJob job = + newBatchJob("batch-job", State.RUNNING, PipelineTranslation.toProto(submitted)); + // The merged job never terminates: waitUntilFinish gives up and returns null. + Mockito.doAnswer( + inv -> { + waitedFor.set(inv.getArgument(0)); + return null; + }) + .when(job) + .waitUntilFinish(any(), any()); + Mockito.doReturn(State.CANCELLED).when(job).cancel(); + batchJob.set(job); + return job; + }); + + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + + // Nobody waited beyond the shortest member timeout; the runaway job was cancelled and every + // member is sent back to be executed again from scratch. + assertThat(rerunReasonOf(f0), containsString("did not terminate within PT3S")); + assertThat(rerunReasonOf(f1), containsString("did not terminate within PT3S")); + assertEquals(Duration.standardSeconds(3), waitedFor.get()); + Mockito.verify(batchJob.get()).cancel(); + assertEquals(1, submittedJobs.get()); + } + + @Test + public void testStalledHarnessFailsTestWithDiagnosticsInsteadOfHanging() throws Exception { + // Short enough to keep the verdict timeout (2 windows + 2 x test timeout = 4 s) small, long + // enough that both members reliably enqueue before the batch would close without the second. + configureBatching(2, 1000L); + Member m0 = new Member("App0").withStep("Step0", 1, 2); + Member m1 = new Member("App1").withStep("Step1", 1, 2); + m0.opts.setTestTimeoutSeconds(1L); + m1.opts.setTestTimeoutSeconds(1L); + + // Simulate the batch runner wedging somewhere it never should (here: inside submission). + CountDownLatch release = new CountDownLatch(1); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + assertThat(invocation.getArgument(0), instanceOf(CompositePipeline.class)); + assertTrue(release.await(WAIT_SECONDS, TimeUnit.SECONDS)); + throw new RuntimeException("released"); + }); + + try { + Future f0 = m0.submit(delegate); + Future f1 = m1.submit(delegate); + for (Future f : Arrays.asList(f0, f1)) { + Throwable cause = causeOf(f); + assertThat(cause, instanceOf(IllegalStateException.class)); + assertThat(cause.getMessage(), containsString("No verdict for test App")); + assertThat(cause.getMessage(), containsString("bug in the test batching machinery")); + // The diagnostics say how far the member got: assembled into a batch, job never submitted. + assertThat(cause.getMessage(), containsString("batch scope: t")); + assertThat(cause.getMessage(), containsString("merged job: never submitted")); + } + } finally { + release.countDown(); + } + } + + @Test + public void testShutdownFailsWaitingMembersPromptlyAndRejectsNewSubmissions() throws Exception { + // A window far longer than the test so the member is parked in the scheduler when we stop it. + configureBatching(2, 60_000L); + Member m0 = new Member("App0").withStep("Step0", 1, 2); + DataflowRunner delegate = Mockito.mock(DataflowRunner.class); + when(delegate.run(any(Pipeline.class))) + .thenThrow(new AssertionError("nothing should be submitted")); + + Future f0 = m0.submit(delegate); + // Let the scheduler take the member into its (open) batch window; whether or not it has by + // the time we stop, the member must fail rather than be stranded. + Thread.sleep(200); + coordinator.shutdown(); + + Throwable cause = causeOf(f0); + assertThat(cause, instanceOf(IllegalStateException.class)); + assertThat(cause.getMessage(), containsString("stopped before scheduling App0")); + + // Later submissions are rejected outright instead of being enqueued onto a dead scheduler. + Member m1 = new Member("App1").withStep("Step1", 1, 2); + Throwable rejected = causeOf(m1.submit(delegate)); + assertThat(rejected, instanceOf(IllegalStateException.class)); + assertThat(rejected.getMessage(), containsString("has been shut down")); + } + + // --------------------------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------------------------- + + /** Returns the cause of the {@link ExecutionException} thrown by {@code future}. */ + private static Throwable causeOf(Future future) { + try { + future.get(WAIT_SECONDS, TimeUnit.SECONDS); + } catch (ExecutionException e) { + return e.getCause(); + } catch (Exception e) { + throw new AssertionError("Unexpected exception", e); + } + fail("Expected the future to fail"); + return null; + } + + /** + * Asserts that the member behind {@code future} was told to have its test executed again as a job + * of its own (see {@link TestPipeline.StandaloneRerunRequested}) and returns the request's + * message, which names the test, the merged job and the reason. + */ + private static String rerunReasonOf(Future future) { + Throwable cause = causeOf(future); + assertThat(cause, instanceOf(TestPipeline.StandaloneRerunRequested.class)); + assertThat(cause.getMessage(), containsString("BeamParallelJunit4Runner")); + return cause.getMessage(); + } + + private DataflowPipelineJob newBatchJob(String jobId, State state, RunnerApi.Pipeline proto) + throws Exception { + DataflowPipelineJob job = + Mockito.spy( + new DataflowPipelineJob(mockClient, jobId, options, Collections.emptyMap(), proto)); + Mockito.doReturn(state).when(job).getState(); + Mockito.doReturn(state).when(job).waitUntilFinish(any(), any()); + return job; + } + + private static DataflowPipelineJob newStandaloneJob(String jobId, State state) throws Exception { + DataflowPipelineJob job = Mockito.mock(DataflowPipelineJob.class); + when(job.getJobId()).thenReturn(jobId); + when(job.getState()).thenReturn(state); + when(job.waitUntilFinish(any(), any())).thenReturn(state); + return job; + } + + /** Returns the id of the transform whose unique name is {@code stepName} or ends in it. */ + private static String transformIdOf(RunnerApi.Pipeline proto, String stepName) { + for (Map.Entry entry : + proto.getComponents().getTransformsMap().entrySet()) { + String uniqueName = entry.getValue().getUniqueName(); + if (uniqueName.equals(stepName) || uniqueName.endsWith("/" + stepName)) { + return entry.getKey(); + } + } + throw new AssertionError("No transform named " + stepName + " in pipeline"); + } + + private static boolean hasTransform(RunnerApi.Pipeline proto, String stepName) { + for (RunnerApi.PTransform t : proto.getComponents().getTransformsMap().values()) { + if (t.getUniqueName().equals(stepName) || t.getUniqueName().endsWith("/" + stepName)) { + return true; + } + } + return false; + } + + /** Returns the batch scope prefix (e.g. {@code t1}) the member owning {@code stepName} got. */ + private static String scopeOf(RunnerApi.Pipeline proto, String stepName) { + String uniqueName = + proto + .getComponents() + .getTransformsMap() + .get(transformIdOf(proto, stepName)) + .getUniqueName(); + int slash = uniqueName.indexOf('/'); + assertTrue("Expected a scoped name but got " + uniqueName, slash > 0); + return uniqueName.substring(0, slash); + } + + private static MetricUpdate passertUpdate( + RunnerApi.Pipeline proto, String stepName, String counterName) { + return createTentativePAssertUpdate(transformIdOf(proto, stepName), counterName, 1); + } + + private static void assertSingleCounter(DataflowPipelineJob job, String step, long expected) { + MetricQueryResults metrics = + job.metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter(MetricNameFilter.named("CustomNs", "myCounter")) + .build()); + List> counters = ImmutableList.copyOf(metrics.getCounters()); + assertEquals(1, counters.size()); + assertEquals( + MetricKey.create(step, MetricName.named("CustomNs", "myCounter")), + counters.get(0).getKey()); + assertThat(counters.get(0).getCommitted(), is(expected)); + } + + private static final class StageInfo { + final String name; + final String state; + final List userNames; + + StageInfo(String name, String state, List userNames) { + this.name = name; + this.state = state; + this.userNames = userNames; + } + } + + private static StageInfo stage(String name, String state, String... userNames) { + return new StageInfo(name, state, Arrays.asList(userNames)); + } + + /** Builds a {@code JOB_VIEW_ALL} job carrying stage states and stage descriptions. */ + private static Job executionDetails(StageInfo... stages) { + StageAttributionBuilder builder = new StageAttributionBuilder(); + for (StageInfo s : stages) { + builder.state(s.name, s.state); + builder.describe(s.name, s.userNames.toArray(new String[0])); + } + return builder.job(); + } + + private static final class StageAttributionBuilder { + private final List states = new ArrayList<>(); + private final List summaries = new ArrayList<>(); + + StageAttributionBuilder state(String stageName, String state) { + states.add( + new ExecutionStageState().setExecutionStageName(stageName).setExecutionStageState(state)); + return this; + } + + StageAttributionBuilder describe(String stageName, String... userNames) { + List transforms = new ArrayList<>(); + for (String userName : userNames) { + transforms.add(new ComponentTransform().setUserName(userName)); + } + summaries.add( + new ExecutionStageSummary().setName(stageName).setComponentTransform(transforms)); + return this; + } + + Job job() { + return new Job() + .setStageStates(states) + .setPipelineDescription(new PipelineDescription().setExecutionPipelineStage(summaries)); + } + + StageAttribution build() { + return StageAttribution.fromJob(job()); + } + } + + private static MetricUpdate createTentativePAssertUpdate( + String stepId, String counterName, int value) { + MetricStructuredName name = + new MetricStructuredName() + .setName(counterName) + .setOrigin("user") + .setContext( + ImmutableMap.of( + "step", stepId, "namespace", PAssert.class.getName(), "tentative", "true")); + return new MetricUpdate().setName(name).setScalar(new BigDecimal(value)); + } + + private static List createCounterUpdates( + String stepId, String namespace, String counterName, long value) { + MetricStructuredName committedName = + new MetricStructuredName() + .setName(counterName) + .setOrigin("user") + .setContext(ImmutableMap.of("step", stepId, "namespace", namespace)); + MetricStructuredName tentativeName = + new MetricStructuredName() + .setName(counterName) + .setOrigin("user") + .setContext( + ImmutableMap.of("step", stepId, "namespace", namespace, "tentative", "true")); + return ImmutableList.of( + new MetricUpdate().setName(committedName).setScalar(new BigDecimal(value)), + new MetricUpdate().setName(tentativeName).setScalar(new BigDecimal(value))); + } +} diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowTestBatchFallbackTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowTestBatchFallbackTest.java new file mode 100644 index 000000000000..f5581e914dfb --- /dev/null +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowTestBatchFallbackTest.java @@ -0,0 +1,132 @@ +/* + * 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.dataflow; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.not; +import static org.junit.Assert.assertEquals; + +import org.apache.beam.runners.dataflow.DataflowTestBatchCoordinator.ScopedDataflowPipelineJob; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.PipelineResult.State; +import org.apache.beam.sdk.testing.BeamParallelJunit4Runner; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.ValidatesRunner; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * ValidatesRunner tests that deliberately exercise {@link DataflowTestBatchCoordinator}'s fallback + * path against the real service: one test fails only while it is part of a merged job, so the + * merged job fails, the coordinator attributes the failure, that test's {@code run()} ends in + * {@link TestPipeline.StandaloneRerunRequested}, and {@link BeamParallelJunit4Runner} executes the + * test again from scratch as a job of its own (where it passes). The other tests are healthy + * companions that give the merged job stages to attribute; they pass whether they are credited from + * the merged job or executed again. + * + *

Keep this class small. Everything merged with the failing test may have to be executed again, + * and since a test JVM runs one test class at a time a merged job only ever contains tests of a + * single class, so the size of this class bounds that cost. All tests run concurrently under {@link + * BeamParallelJunit4Runner} so that they do end up in the same batch. + * + *

When batching is disabled (streaming tasks, {@code -PtestBatching=false}) every test here runs + * standalone and simply passes. + */ +@RunWith(BeamParallelJunit4Runner.class) +public class DataflowTestBatchFallbackTest { + private static final Logger LOG = LoggerFactory.getLogger(DataflowTestBatchFallbackTest.class); + + @Rule public final transient TestPipeline p = TestPipeline.create(); + + /** Fails inside a merged job (recognised by its job name) and passes in a standalone one. */ + private static class FailWhenMergedFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) { + String jobName = c.getPipelineOptions().getJobName(); + if (jobName != null + && jobName.startsWith(DataflowTestBatchCoordinator.BATCH_JOB_NAME_PREFIX)) { + throw new IllegalStateException( + "Deliberate failure inside merged Dataflow test job " + jobName); + } + c.output(c.element()); + } + } + + private static class IdentityFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) { + c.output(c.element()); + } + } + + @Test + @Category(ValidatesRunner.class) + public void testMemberFailingOnlyWhenMergedFallsBackAndPasses() { + PCollection out = + p.apply(Create.of(1, 2, 3)).apply("FailWhenMerged", ParDo.of(new FailWhenMergedFn())); + PAssert.that(out).containsInAnyOrder(1, 2, 3); + + PipelineResult result = p.run(); + + // This test can never be credited from a merged job: its DoFn throws before any PAssert runs + // there, so the merged attempt has no success counters for it. Either batching was off, or this + // is the standalone re-execution of the test. + assertThat(result, not(instanceOf(ScopedDataflowPipelineJob.class))); + assertEquals(State.DONE, result.getState()); + } + + @Test + @Category(ValidatesRunner.class) + public void testHealthyCompanionA() { + runHealthyCompanion("CompanionA"); + } + + @Test + @Category(ValidatesRunner.class) + public void testHealthyCompanionB() { + runHealthyCompanion("CompanionB"); + } + + private void runHealthyCompanion(String name) { + PCollection out = p.apply(Create.of(4, 5, 6)).apply(name, ParDo.of(new IdentityFn())); + PAssert.that(out).containsInAnyOrder(4, 5, 6); + + PipelineResult result = p.run(); + + // Whether the coordinator could credit this test from the failed merged job (stage states + // attributable) or had to re-run it is informative for the postcommit log, but both are passes. + LOG.info( + "{} finished {} ({})", + name, + result.getState(), + result instanceof ScopedDataflowPipelineJob + ? "credited from merged job" + : "standalone or re-run"); + assertEquals(State.DONE, result.getState()); + } +} diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java index 5cd16757e36b..f22766043998 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java @@ -19,7 +19,10 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.empty; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.notNullValue; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; @@ -32,12 +35,18 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.api.client.googleapis.json.GoogleJsonError; +import com.google.api.client.googleapis.json.GoogleJsonResponseException; +import com.google.api.client.http.HttpHeaders; +import com.google.api.client.http.HttpResponseException; +import com.google.api.client.util.Sleeper; import com.google.api.services.dataflow.model.JobMessage; import com.google.api.services.dataflow.model.JobMetrics; import com.google.api.services.dataflow.model.MetricStructuredName; import com.google.api.services.dataflow.model.MetricUpdate; import java.io.IOException; import java.math.BigDecimal; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -669,4 +678,106 @@ public boolean matches(Object o) { @Override public void describeTo(Description description) {} } + + /** A sleeper that records requested sleeps instead of performing them. */ + private static class RecordingSleeper implements Sleeper { + final List sleepsMillis = new ArrayList<>(); + + @Override + public void sleep(long millis) { + sleepsMillis.add(millis); + } + } + + private static RuntimeException quotaRejection() { + GoogleJsonError details = new GoogleJsonError(); + details.setMessage("Quota exceeded for quota metric 'Job creation requests'"); + GoogleJsonResponseException http429 = + new GoogleJsonResponseException( + new HttpResponseException.Builder(429, "Too Many Requests", new HttpHeaders()), + details); + // DataflowRunner wraps the service exception exactly like this. + return new RuntimeException( + "Failed to create a workflow job: " + details.getMessage(), http429); + } + + @Test + public void testSubmitRetriesQuotaRejections() { + Pipeline p = Pipeline.create(options); + DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class); + DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class); + when(mockRunner.run(any(Pipeline.class))) + .thenThrow(quotaRejection()) + .thenThrow(quotaRejection()) + .thenReturn(mockJob); + + TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); + RecordingSleeper sleeper = new RecordingSleeper(); + runner.setSleeper(sleeper); + + assertSame(mockJob, runner.submit(mockRunner, p)); + verify(mockRunner, Mockito.times(3)).run(p); + assertEquals(2, sleeper.sleepsMillis.size()); + for (long sleep : sleeper.sleepsMillis) { + assertThat(sleep, greaterThan(0L)); + } + } + + @Test + public void testSubmitDoesNotRetryOtherFailures() { + Pipeline p = Pipeline.create(options); + DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class); + RuntimeException failure = new RuntimeException("Failed to create a workflow job: bad graph"); + when(mockRunner.run(any(Pipeline.class))).thenThrow(failure); + + TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); + RecordingSleeper sleeper = new RecordingSleeper(); + runner.setSleeper(sleeper); + + try { + runner.submit(mockRunner, p); + fail("Expected the submission failure to propagate"); + } catch (RuntimeException e) { + assertSame(failure, e); + } + verify(mockRunner, Mockito.times(1)).run(p); + assertThat(sleeper.sleepsMillis, empty()); + } + + @Test + public void testSubmitGivesUpAfterRepeatedQuotaRejections() { + Pipeline p = Pipeline.create(options); + DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class); + when(mockRunner.run(any(Pipeline.class))) + .thenAnswer( + invocation -> { + throw quotaRejection(); + }); + + TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); + RecordingSleeper sleeper = new RecordingSleeper(); + runner.setSleeper(sleeper); + + try { + runner.submit(mockRunner, p); + fail("Expected the submission to fail once retries are exhausted"); + } catch (RuntimeException e) { + assertThat(TestDataflowRunner.quotaRejection(e), notNullValue()); + } + // One initial attempt plus a bounded number of retries, each preceded by a sleep. + int attempts = Mockito.mockingDetails(mockRunner).getInvocations().size(); + assertThat(attempts, greaterThan(1)); + assertEquals(attempts - 1, sleeper.sleepsMillis.size()); + } + + @Test + public void testQuotaRejectionDetection() { + assertThat(TestDataflowRunner.quotaRejection(quotaRejection()), notNullValue()); + assertNull(TestDataflowRunner.quotaRejection(new RuntimeException("boom"))); + GoogleJsonResponseException http400 = + new GoogleJsonResponseException( + new HttpResponseException.Builder(400, "Bad Request", new HttpHeaders()), + new GoogleJsonError()); + assertNull(TestDataflowRunner.quotaRejection(new RuntimeException("wrapped", http400))); + } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/BeamParallelJunit4Runner.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/BeamParallelJunit4Runner.java index 59527b161f08..6df2c8cf36c8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/BeamParallelJunit4Runner.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/BeamParallelJunit4Runner.java @@ -17,23 +17,32 @@ */ package org.apache.beam.sdk.testing; +import java.lang.annotation.Annotation; import java.lang.annotation.ElementType; import java.lang.annotation.Retention; import java.lang.annotation.RetentionPolicy; import java.lang.annotation.Target; +import java.util.Collection; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicInteger; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.experimental.categories.Category; +import org.junit.runner.Description; import org.junit.runner.notification.RunNotifier; import org.junit.runners.BlockJUnit4ClassRunner; import org.junit.runners.model.FrameworkMethod; import org.junit.runners.model.InitializationError; +import org.junit.runners.model.MultipleFailureException; import org.junit.runners.model.RunnerScheduler; +import org.junit.runners.model.Statement; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A JUnit 4 runner that supports concurrent execution of {@code @Test} methods within a test class, @@ -49,10 +58,18 @@ * *

Classes or methods annotated with {@link SerialTest} are always executed sequentially on the * calling thread after draining any in-flight parallel test methods in the class. + * + *

Runners that merge the pipelines of concurrently running tests into one job rely on this + * runner in two ways: it lets each {@link TestPipeline} rule see its test instance (to recognize + * tests that must not be merged), and it re-executes a test once, from scratch and with merging + * disabled, when the runner could not reach a verdict for it from a merged job (see {@link + * TestPipeline.StandaloneRerunRequested}). */ @Internal public final class BeamParallelJunit4Runner extends BlockJUnit4ClassRunner { + private static final Logger LOG = LoggerFactory.getLogger(BeamParallelJunit4Runner.class); + /** * Marks a test class or {@code @Test} method as requiring serial (non-parallel) execution even * when parallel test execution is enabled. @@ -108,6 +125,18 @@ public void finished() { }); } + /** + * Hands the test instance to its {@link TestPipeline} rule(s), which inspect the instance's other + * rules to decide whether the test's pipeline may be merged with others. A {@code TestRule} never + * sees the instance itself; the runner is the only place that does. + */ + @Override + protected Object createTest() throws Exception { + Object instance = super.createTest(); + TestPipeline.attachTestInstance(instance); + return instance; + } + private void awaitPendingFutures() { try { pendingFutures.join(); @@ -116,6 +145,91 @@ private void awaitPendingFutures() { } } + // --------------------------------------------------------------------------------------------- + // Re-executing a test whose pipeline was merged into a shared job without a verdict + // --------------------------------------------------------------------------------------------- + // + // A runner that merges several tests' pipelines into one job may be unable to tell from that job + // whether a given test passed (see TestPipeline.StandaloneRerunRequested). The pipeline object + // cannot simply be run again, so the test is executed a second time from scratch: new instance, + // @Before/@After, rules and all, with its TestPipeline told to insist on a job of its own. Both + // attempts are one test as far as JUnit is concerned; it sees a single start/finish pair. + + /** + * Tests currently in their standalone re-execution, by {@link Description}. Keyed by description + * rather than thread because the test body does not necessarily run on the thread that started it + * (JUnit's Timeout rule, for one, evaluates it on a thread of its own). A method is never + * re-executed concurrently with itself, so an entry unambiguously refers to the current attempt. + */ + private static final Set STANDALONE_RERUNS = ConcurrentHashMap.newKeySet(); + + /** + * Returns {@code true} while the test described by {@code description} is being re-executed after + * an inconclusive merged attempt; {@link TestPipeline} consults this when its rule is applied. + */ + static boolean isStandaloneRerun(Description description) { + return STANDALONE_RERUNS.contains(description); + } + + /** + * The {@link TestPipeline.StandaloneRerunRequested} behind {@code failure}, or {@code null} if + * there is none. Test code and JUnit itself may wrap what {@code run()} threw ({@code + * assertThrows}, {@code @Test(expected)}, an {@code @After} failing alongside it), so the cause + * chain and {@link MultipleFailureException}'s failures are searched too. + */ + @VisibleForTesting + static TestPipeline.@Nullable StandaloneRerunRequested standaloneRerunRequest(Throwable failure) { + int depth = 0; + for (Throwable t = failure; t != null && depth < 32; t = t.getCause(), depth++) { + if (t instanceof TestPipeline.StandaloneRerunRequested) { + return (TestPipeline.StandaloneRerunRequested) t; + } + if (t instanceof MultipleFailureException) { + for (Throwable each : ((MultipleFailureException) t).getFailures()) { + TestPipeline.StandaloneRerunRequested nested = standaloneRerunRequest(each); + if (nested != null) { + return nested; + } + } + } + } + return null; + } + + /** + * Like {@link BlockJUnit4ClassRunner#runChild}, except that a test ending in {@link + * TestPipeline.StandaloneRerunRequested} is executed a second time (see above). A request raised + * by that second attempt is a bug in the requesting runner and is reported as the test's failure. + */ + private void runChildInternal(final FrameworkMethod method, final RunNotifier notifier) { + final Description description = describeChild(method); + Statement statement = + new Statement() { + @Override + public void evaluate() throws Throwable { + try { + methodBlock(method).evaluate(); + } catch (Throwable firstAttempt) { + TestPipeline.StandaloneRerunRequested request = standaloneRerunRequest(firstAttempt); + if (request == null) { + throw firstAttempt; + } + LOG.info( + "Re-executing {} with its pipeline as a job of its own. {}", + description.getDisplayName(), + String.valueOf(request)); + STANDALONE_RERUNS.add(description); + try { + methodBlock(method).evaluate(); + } finally { + STANDALONE_RERUNS.remove(description); + } + } + } + }; + runLeaf(statement, description, notifier); + } + @Override protected void runChild(final FrameworkMethod method, final RunNotifier notifier) { if (isIgnored(method)) { @@ -124,7 +238,7 @@ protected void runChild(final FrameworkMethod method, final RunNotifier notifier } ExecutorService executor = selectExecutor(method); if (executor == null) { - super.runChild(method, notifier); + runChildInternal(method, notifier); return; } if (isClassMarkedSerial(getTestClass().getJavaClass()) || isMethodMarkedSerial(method)) { @@ -132,13 +246,13 @@ protected void runChild(final FrameworkMethod method, final RunNotifier notifier // this test // serially. awaitPendingFutures(); - super.runChild(method, notifier); + runChildInternal(method, notifier); return; } pendingFutures = CompletableFuture.allOf( pendingFutures, - CompletableFuture.runAsync(() -> super.runChild(method, notifier), executor)); + CompletableFuture.runAsync(() -> runChildInternal(method, notifier), executor)); } private @Nullable ExecutorService selectExecutor(FrameworkMethod method) { @@ -171,7 +285,21 @@ private static boolean isMethodMarkedSerial(FrameworkMethod method) { return method.getAnnotation(SerialTest.class) != null; } - private static boolean isClassMarkedSerial(@Nullable Class clazz) { + /** Returns {@code true} if {@code annotations} contains {@link SerialTest}. */ + static boolean hasSerialAnnotation(Collection annotations) { + for (Annotation annotation : annotations) { + if (annotation.annotationType() == SerialTest.class) { + return true; + } + } + return false; + } + + /** + * Returns {@code true} if {@code clazz}, any of its superclasses, or any of its enclosing classes + * is annotated with {@link SerialTest}. + */ + static boolean isClassMarkedSerial(@Nullable Class clazz) { if (clazz == null || clazz == Object.class) { return false; } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java index 4406bf0d7692..fc34401c9c50 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java @@ -75,6 +75,7 @@ import org.apache.beam.sdk.values.PDone; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; @@ -123,10 +124,13 @@ public class PAssert { Metrics.counter(PAssert.class, PAssert.FAILURE_COUNTER); // Atomic so that PAssert transforms constructed concurrently (e.g. by tests running under - // BeamParallelJunit4Runner) never receive duplicate names within a pipeline. + // BeamParallelJunit4Runner) never receive duplicate names within a pipeline. With a plain int a + // lost update can move the counter backwards, so a single thread can be handed a name it already + // used in its own pipeline, which stableUniqueNames=ERROR then rejects. private static final AtomicInteger assertCount = new AtomicInteger(0); - private static String nextAssertionName() { + @VisibleForTesting + static String nextAssertionName() { return "PAssert$" + assertCount.getAndIncrement(); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestPipeline.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestPipeline.java index b157e4ad785d..76f6c69cc01d 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestPipeline.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestPipeline.java @@ -25,11 +25,17 @@ import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.lang.annotation.Annotation; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.annotations.Internal; @@ -46,7 +52,10 @@ import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; import org.apache.beam.sdk.runners.TransformHierarchy; import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.transforms.resourcehints.ResourceHints; import org.apache.beam.sdk.util.common.ReflectHelpers; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Optional; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicate; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates; @@ -56,10 +65,14 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; import org.junit.experimental.categories.Category; +import org.junit.rules.ExpectedException; import org.junit.rules.TestRule; import org.junit.runner.Description; import org.junit.runners.model.Statement; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A creator of test pipelines that can be used inside of tests that can be configured to run @@ -130,6 +143,7 @@ * documentation section. */ public class TestPipeline extends Pipeline implements TestRule { + private static final Logger LOG = LoggerFactory.getLogger(TestPipeline.class); private final PipelineOptions options; @@ -306,11 +320,12 @@ static TestPipeline createWithEnforcement() { } public static TestPipeline fromOptions(PipelineOptions options) { - return new TestPipeline(options); + return new TestPipeline(new TransformHierarchy(ResourceHints.fromOptions(options)), options); } - private TestPipeline(final PipelineOptions options) { - super(options); + private TestPipeline(final TransformHierarchy hierarchy, final PipelineOptions options) { + super(hierarchy, options); + this.hierarchy = hierarchy; this.options = options; } @@ -324,10 +339,7 @@ void setDeducedEnforcementLevel(Collection annotations) { // if the enforcement level has not been set by the user do auto-inference if (!enforcement.isPresent()) { - final boolean annotatedWithNeedsRunner = - FluentIterable.from(annotations) - .filter(Annotations.Predicates.isAnnotationOfType(Category.class)) - .anyMatch(Annotations.Predicates.isCategoryOf(NeedsRunner.class, true)); + final boolean annotatedWithNeedsRunner = hasCategory(annotations, NeedsRunner.class); final boolean crashingRunner = CrashingRunner.class.isAssignableFrom(options.getRunner()); @@ -348,15 +360,56 @@ void afterUserCodeFinished() { enforcement.get().afterUserCodeFinished(); } + private boolean standaloneExecutionRequired = false; + private int initialOptionsRevision = -1; + + private static boolean hasCategory(Collection annotations, Class targetCategory) { + return FluentIterable.from(annotations) + .filter(Annotations.Predicates.isAnnotationOfType(Category.class)) + .anyMatch(Annotations.Predicates.isCategoryOf(targetCategory, true)); + } + + /** + * Whether {@code testClass}, a superclass, or an enclosing class carries a {@link Category} that + * is or extends {@code targetCategory}. A method's {@link Description} only exposes the method's + * own annotations, so class-level categories have to be looked up here. + */ + private static boolean hasClassCategory(@Nullable Class testClass, Class targetCategory) { + for (Class clazz = testClass; clazz != null; clazz = clazz.getEnclosingClass()) { + Category category = clazz.getAnnotation(Category.class); // @Inherited: covers superclasses + if (category != null + && Annotations.Predicates.isCategoryOf(targetCategory, true).apply(category)) { + return true; + } + } + return false; + } + @Override public Statement apply(final Statement statement, final Description description) { return new Statement() { @Override public void evaluate() throws Throwable { - options.as(ApplicationNameOptions.class).setAppName(getAppName(description)); + String appName = getAppName(description); + options.as(ApplicationNameOptions.class).setAppName(appName); + initialOptionsRevision = options.revision(); + if (Boolean.getBoolean(PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES)) { + nameRoot(appName); + } - setDeducedEnforcementLevel(description.getAnnotations()); + Collection annotations = description.getAnnotations(); + setDeducedEnforcementLevel(annotations); + + Test testAnnotation = description.getAnnotation(Test.class); + if ((testAnnotation != null && testAnnotation.expected() != Test.None.class) + || hasCategory(annotations, UsesFailureMessage.class) + || hasClassCategory(description.getTestClass(), UsesFailureMessage.class) + || BeamParallelJunit4Runner.hasSerialAnnotation(annotations) + || BeamParallelJunit4Runner.isClassMarkedSerial(description.getTestClass()) + || BeamParallelJunit4Runner.isStandaloneRerun(description)) { + standaloneExecutionRequired = true; + } // statement.evaluate() essentially runs the user code contained in the unit test at hand. // Exceptions thrown during the execution of the user's test code will propagate here, @@ -372,6 +425,268 @@ public void evaluate() throws Throwable { }; } + // --------------------------------------------------------------------------------------------- + // Unique root names + // --------------------------------------------------------------------------------------------- + + /** + * System property which, when {@code true}, gives every {@link TestPipeline} a unique root + * transform name of the form {@code t-} as soon as its JUnit rule is applied. Every + * transform and {@link PCollection} name in the test is then prefixed by it (for example {@code + * t12-ParDoTest$BasicTests-testParDo/ParDo(Anonymous)}), so the graphs of different tests never + * share a name. Runners that merge several test pipelines into one job rely on this; see {@link + * #getRootName()}. + */ + public static final String PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES = + "beamTestPipelineUniqueRootNames"; + + /** + * Upper bound on a root name. Generous enough that real test names ({@code Class$Nested-method}, + * typically 40-80 characters) survive intact, since the method name at the end is the most + * informative part, while still bounding pathological names. + */ + @VisibleForTesting static final int MAX_ROOT_NAME_LENGTH = 100; + + private static final AtomicInteger ROOT_NAME_COUNTER = new AtomicInteger(1); + + private final TransformHierarchy hierarchy; + private String rootName = ""; + + /** + * For internal use only; no backwards-compatibility guarantees. + * + *

Returns the unique root transform name of this pipeline, or the empty string if it has none + * (see {@link #PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES}). When non-empty, every full + * transform name in this pipeline starts with {@code getRootName() + "/"}. + */ + @Internal + public String getRootName() { + return rootName; + } + + /** + * For internal use only; no backwards-compatibility guarantees. + * + *

Gives this pipeline a unique root name of the form {@code t-} (sanitized and + * truncated). Normally invoked by the JUnit rule or extension; only has an effect if no transform + * has been applied yet, otherwise a warning is logged and the pipeline stays unnamed. + */ + @Internal + public void nameRoot(String testName) { + String candidate = rootNameFor(ROOT_NAME_COUNTER.getAndIncrement(), testName); + try { + hierarchy.setRootName(candidate); + rootName = candidate; + } catch (IllegalStateException e) { + LOG.warn( + "Not giving the TestPipeline of {} a unique root name: transforms were applied to it" + + " before the test started. Its transform names stay unprefixed.", + testName); + } + } + + /** + * Builds {@code t-}. The sequence number alone makes the name unique, so the + * test name is only a readability aid: it is sanitized to a conservative character set (notably + * no {@code /}, which separates name segments) and truncated to keep the root within {@link + * #MAX_ROOT_NAME_LENGTH}. + */ + @VisibleForTesting + static String rootNameFor(int sequence, String testName) { + String prefix = "t" + sequence + "-"; + String hint = testName.replaceAll("[^A-Za-z0-9_.$-]", "_"); + int room = MAX_ROOT_NAME_LENGTH - prefix.length(); + if (hint.length() > room) { + hint = hint.substring(0, Math.max(0, room)); + } + return prefix + hint; + } + + /** + * For internal use only; no backwards-compatibility guarantees. + * + *

Returns {@code true} if this {@link TestPipeline} should not be merged into a shared batch + * pipeline (for example, when the test expects an exception or assertion failure, uses custom + * per-test option arguments, is marked for serial execution, or is being re-executed after a + * merged attempt; see {@link StandaloneRerunRequested}). + */ + @Internal + public boolean isStandaloneExecutionRequired() { + return standaloneExecutionRequired + || (initialOptionsRevision >= 0 && options.revision() != initialOptionsRevision) + || !providerRuntimeValues.isEmpty() + || testExpectsException(); + } + + /** + * For internal use only; no backwards-compatibility guarantees. + * + *

Thrown out of {@link #run()} by a runner that merged this pipeline into a job shared with + * other tests and could not reach a verdict for this test from that job. The pipeline is not run + * again by the runner: once the shared job was built from it, this pipeline's graph is no longer + * guaranteed to be the one the test constructed. Instead the whole test has to be executed again + * from scratch, with {@linkplain #isStandaloneExecutionRequired() standalone execution} forced so + * that the fresh pipeline runs as a job of its own. {@link BeamParallelJunit4Runner} does that + * automatically (once); under any other runner this surfaces as the test's failure. + * + *

An {@link Error} rather than an exception so that test code wrapping {@code run()} in {@code + * catch (Exception e)} or {@code assertThrows(Exception.class, ...)} cannot mistake it for the + * failure it was looking for. + */ + @Internal + public static final class StandaloneRerunRequested extends Error { + public StandaloneRerunRequested(String message) { + super(message); + } + } + + // --------------------------------------------------------------------------------------------- + // Detecting tests that expect run() to fail + // --------------------------------------------------------------------------------------------- + // + // A test that expects run() to throw must not be merged with other tests: it would fail the + // shared job. Static signals (@Test(expected), categories, serial markers) are read from the + // JUnit Description in apply(). The dynamic one, an ExpectedException rule on which the test + // has called expect(...), lives on the test instance, which a TestRule never sees; the runner + // hands the instance to its TestPipeline rule(s) through attachTestInstance so that it can be + // inspected here at run() time. The instance is attached to the rule rather than published in a + // thread-local because the test body does not necessarily run on the thread that created the + // instance: JUnit's Timeout rule, for one, evaluates the test on a thread of its own. + // A missed detection is a performance problem (that test and its co-members re-run standalone), + // never a correctness one. + + /** The test-class instance this rule belongs to, if the runner has told us; see above. */ + private transient @Nullable Object testInstance; + + /** + * {@code ExpectedException#isAnyExceptionExpected()}, the rule's own (private) notion of whether + * the test has called {@code expect(...)}. Resolved once; {@code null} if this JUnit version does + * not have it, in which case such tests are simply not recognized. + */ + private static final @Nullable Method IS_ANY_EXCEPTION_EXPECTED = lookupIsAnyExceptionExpected(); + + private static final AtomicBoolean PROBE_FAILURE_LOGGED = new AtomicBoolean(false); + + /** {@link ExpectedException}-typed fields per test class, resolved once. */ + private static final ConcurrentHashMap, List> EXPECTED_EXCEPTION_FIELDS = + new ConcurrentHashMap<>(); + + /** {@link TestPipeline}-typed fields per test class, resolved once. */ + private static final ConcurrentHashMap, List> TEST_PIPELINE_FIELDS = + new ConcurrentHashMap<>(); + + /** + * Attaches {@code testInstance} to every {@link TestPipeline} held in one of its fields (declared + * in its class or a superclass), so that those pipelines can inspect the instance's other rules. + * Called by {@link BeamParallelJunit4Runner} for each freshly created test instance. + */ + static void attachTestInstance(Object testInstance) { + for (Field field : + fieldsOfType(TEST_PIPELINE_FIELDS, testInstance.getClass(), TestPipeline.class)) { + try { + Object value = field.get(testInstance); + if (value instanceof TestPipeline) { + ((TestPipeline) value).testInstance = testInstance; + } + } catch (ReflectiveOperationException | RuntimeException e) { + if (PROBE_FAILURE_LOGGED.compareAndSet(false, true)) { + LOG.warn( + "Failed to read TestPipeline field {} of {}; tests that expect an exception via an" + + " ExpectedException rule may be merged into shared test jobs.", + field.getName(), + testInstance.getClass().getName(), + e); + } + } + } + } + + /** + * Returns {@code true} if the test this pipeline belongs to has armed an {@link + * ExpectedException} rule, i.e. it expects to fail. Only {@link ExpectedException} fields of the + * test instance are consulted; tests that expect failure some other way (for example {@code + * assertThrows} around {@code run()}) are not recognized, and neither is anything when the test + * runner has not attached the instance (see {@link #attachTestInstance}). + */ + @VisibleForTesting + boolean testExpectsException() { + Object instance = testInstance; + if (instance == null || IS_ANY_EXCEPTION_EXPECTED == null) { + return false; + } + for (Field field : + fieldsOfType(EXPECTED_EXCEPTION_FIELDS, instance.getClass(), ExpectedException.class)) { + try { + Object rule = field.get(instance); + if (rule instanceof ExpectedException + && Boolean.TRUE.equals(IS_ANY_EXCEPTION_EXPECTED.invoke(rule))) { + return true; + } + } catch (ReflectiveOperationException | RuntimeException e) { + if (PROBE_FAILURE_LOGGED.compareAndSet(false, true)) { + LOG.warn( + "Failed to inspect ExpectedException rule {} of {}; tests that expect an exception" + + " via that rule may be merged into shared test jobs.", + field.getName(), + instance.getClass().getName(), + e); + } + } + } + return false; + } + + private static @Nullable Method lookupIsAnyExceptionExpected() { + try { + Method method = ExpectedException.class.getDeclaredMethod("isAnyExceptionExpected"); + method.setAccessible(true); + return method; + } catch (ReflectiveOperationException | RuntimeException e) { + LOG.warn( + "Cannot inspect JUnit's ExpectedException rule; tests that expect an exception via that" + + " rule will not be recognized as such and may be merged into shared test jobs.", + e); + return null; + } + } + + /** + * The instance fields of {@code testClass} (or a superclass) whose declared type is assignable to + * {@code type}, made accessible, cached in {@code cache}. + */ + private static List fieldsOfType( + ConcurrentHashMap, List> cache, Class testClass, Class type) { + return cache.computeIfAbsent( + testClass, + k -> { + List fields = new ArrayList<>(); + for (Class clazz = k; + clazz != null && clazz != Object.class; + clazz = clazz.getSuperclass()) { + for (Field field : clazz.getDeclaredFields()) { + if (Modifier.isStatic(field.getModifiers()) + || !type.isAssignableFrom(field.getType())) { + continue; + } + try { + field.setAccessible(true); + fields.add(field); + } catch (RuntimeException e) { + if (PROBE_FAILURE_LOGGED.compareAndSet(false, true)) { + LOG.warn( + "Cannot access field {} of {}; tests that expect an exception via an" + + " ExpectedException rule may be merged into shared test jobs.", + field.getName(), + k.getName(), + e); + } + } + } + } + return fields; + }); + } + /** * Runs this {@link TestPipeline}, unwrapping any {@code AssertionError} that is raised during * testing. @@ -390,6 +705,7 @@ public PipelineResult run() { *

Most of logic is similar to {@link #testingPipelineOptions}. */ public PipelineResult runWithAdditionalOptionArgs(List additionalArgs) { + standaloneExecutionRequired = true; try { String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS, ""); List args = new ArrayList<>(); @@ -430,6 +746,10 @@ public PipelineResult run(PipelineOptions options) { enforcement.isPresent(), "Is your TestPipeline declaration missing a @Rule annotation? Usage: " + "@Rule public final transient TestPipeline pipeline = TestPipeline.create();"); + if (options != this.options + || (initialOptionsRevision >= 0 && options.revision() != initialOptionsRevision)) { + standaloneExecutionRequired = true; + } final PipelineResult pipelineResult; try { diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/BeamParallelJunit4RunnerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/BeamParallelJunit4RunnerTest.java index 1fd2b168bd06..7f16bb5266c3 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/BeamParallelJunit4RunnerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/BeamParallelJunit4RunnerTest.java @@ -17,10 +17,20 @@ */ package org.apache.beam.sdk.testing; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.not; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeTrue; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; @@ -28,14 +38,20 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import org.apache.beam.sdk.testing.BeamParallelJunit4Runner.SerialTest; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; +import org.junit.rules.ExpectedException; import org.junit.rules.TestRule; +import org.junit.rules.Timeout; import org.junit.runner.JUnitCore; import org.junit.runner.Result; import org.junit.runner.RunWith; +import org.junit.runner.notification.Failure; import org.junit.runners.JUnit4; +import org.junit.runners.model.MultipleFailureException; /** Tests for {@link BeamParallelJunit4Runner}. */ @RunWith(JUnit4.class) @@ -200,4 +216,142 @@ public void testSerialTestMethodDrainsAndRunsAlone() { IN_TEST_HARNESS.set(false); } } + + @RunWith(BeamParallelJunit4Runner.class) + public static class SampleExpectedExceptionCases { + static final AtomicBoolean OBSERVED_WHEN_EXPECTING = new AtomicBoolean(false); + static final AtomicBoolean OBSERVED_WHEN_NOT_EXPECTING = new AtomicBoolean(true); + static final AtomicBoolean OBSERVED_ON_CHILD_THREAD = new AtomicBoolean(false); + + @Rule public ExpectedException thrown = ExpectedException.none(); + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @Test + public void testExpecting() throws InterruptedException { + assumeTrue(IN_TEST_HARNESS.get()); + thrown.expect(IllegalStateException.class); + OBSERVED_WHEN_EXPECTING.set(pipeline.testExpectsException()); + // Detection is tied to the test instance, not to a thread: rules such as Timeout evaluate + // the test body on a thread of their own. + Thread child = + new Thread(() -> OBSERVED_ON_CHILD_THREAD.set(pipeline.testExpectsException())); + child.start(); + child.join(); + throw new IllegalStateException("expected"); + } + + @Test + public void testNotExpecting() { + assumeTrue(IN_TEST_HARNESS.get()); + OBSERVED_WHEN_NOT_EXPECTING.set(pipeline.testExpectsException()); + } + } + + @Test + public void testDetectsActiveExpectedExceptionRule() { + IN_TEST_HARNESS.set(true); + try { + Result result = JUnitCore.runClasses(SampleExpectedExceptionCases.class); + assertEquals(0, result.getFailureCount()); + assertEquals(2, result.getRunCount()); + assertTrue(SampleExpectedExceptionCases.OBSERVED_WHEN_EXPECTING.get()); + assertFalse(SampleExpectedExceptionCases.OBSERVED_WHEN_NOT_EXPECTING.get()); + assertTrue(SampleExpectedExceptionCases.OBSERVED_ON_CHILD_THREAD.get()); + } finally { + IN_TEST_HARNESS.set(false); + } + } + + @RunWith(BeamParallelJunit4Runner.class) + public static class SampleStandaloneRerunCases { + static final AtomicInteger BEFORE_CALLS = new AtomicInteger(0); + static final List RERUN_INSTANCES = Collections.synchronizedList(new ArrayList<>()); + static final AtomicBoolean STANDALONE_ON_FIRST_ATTEMPT = new AtomicBoolean(true); + static final AtomicBoolean STANDALONE_ON_RERUN = new AtomicBoolean(false); + static final AtomicInteger ALWAYS_REQUESTING_ATTEMPTS = new AtomicInteger(0); + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + // Evaluates the test body on a thread of its own, as in the ValidatesRunner suites. + @Rule public final Timeout timeout = Timeout.seconds(30); + + @Before + public void before() { + BEFORE_CALLS.incrementAndGet(); + } + + @Test + public void testRerunOnce() { + assumeTrue(IN_TEST_HARNESS.get()); + RERUN_INSTANCES.add(this); + if (RERUN_INSTANCES.size() == 1) { + STANDALONE_ON_FIRST_ATTEMPT.set(pipeline.isStandaloneExecutionRequired()); + // Wrapped, as test code catching what run() threw would do. + throw new RuntimeException( + "wrapped", new TestPipeline.StandaloneRerunRequested("no verdict from merged job")); + } + STANDALONE_ON_RERUN.set(pipeline.isStandaloneExecutionRequired()); + } + + @Test + public void testAlwaysRequesting() { + assumeTrue(IN_TEST_HARNESS.get()); + ALWAYS_REQUESTING_ATTEMPTS.incrementAndGet(); + throw new TestPipeline.StandaloneRerunRequested("still no verdict"); + } + } + + @Test + public void testReexecutesTestOnceWhenStandaloneRerunIsRequested() { + System.setProperty(BeamParallelJunit4Runner.DEFAULT_TEST_THREADS_PROPERTY, "2"); + IN_TEST_HARNESS.set(true); + try { + Result result = JUnitCore.runClasses(SampleStandaloneRerunCases.class); + // Two attempts are one test to JUnit. + assertEquals(2, result.getRunCount()); + + // testRerunOnce: a second, fresh instance (and @Before) with standalone execution forced. + assertEquals(2, SampleStandaloneRerunCases.RERUN_INSTANCES.size()); + assertNotSame( + SampleStandaloneRerunCases.RERUN_INSTANCES.get(0), + SampleStandaloneRerunCases.RERUN_INSTANCES.get(1)); + assertFalse(SampleStandaloneRerunCases.STANDALONE_ON_FIRST_ATTEMPT.get()); + assertTrue(SampleStandaloneRerunCases.STANDALONE_ON_RERUN.get()); + + // testAlwaysRequesting: re-executed exactly once, then the request is the failure. + assertEquals(2, SampleStandaloneRerunCases.ALWAYS_REQUESTING_ATTEMPTS.get()); + assertEquals(1, result.getFailureCount()); + Failure failure = result.getFailures().get(0); + assertEquals("testAlwaysRequesting", failure.getDescription().getMethodName()); + assertThat(failure.getException(), instanceOf(TestPipeline.StandaloneRerunRequested.class)); + + // Two tests, two attempts each. + assertEquals(4, SampleStandaloneRerunCases.BEFORE_CALLS.get()); + } finally { + IN_TEST_HARNESS.set(false); + } + } + + @Test + public void testStandaloneRerunRequestIsFoundBehindWrappers() { + TestPipeline.StandaloneRerunRequested request = + new TestPipeline.StandaloneRerunRequested("no verdict"); + assertSame(request, BeamParallelJunit4Runner.standaloneRerunRequest(request)); + assertSame( + request, + BeamParallelJunit4Runner.standaloneRerunRequest( + new AssertionError("unexpected exception type", request))); + assertSame( + request, + BeamParallelJunit4Runner.standaloneRerunRequest( + new MultipleFailureException( + ImmutableList.of(new IllegalStateException("@After failed"), request)))); + assertNull(BeamParallelJunit4Runner.standaloneRerunRequest(new RuntimeException("other"))); + // A cause cycle must not hang detection. + RuntimeException a = new RuntimeException("a"); + RuntimeException b = new RuntimeException("b", a); + a.initCause(b); + assertNull(BeamParallelJunit4Runner.standaloneRerunRequest(a)); + // It is an Error so that test code catching Exception around run() cannot swallow it. + assertThat(request, not(instanceOf(Exception.class))); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java index bec3de5a6241..21b746dccd4a 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java @@ -31,7 +31,15 @@ import java.io.InputStream; import java.io.OutputStream; import java.io.Serializable; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.regex.Pattern; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.AtomicCoder; @@ -138,6 +146,37 @@ private void throwWrappedError() { } } + @Test + public void testAssertionNamesAreUniqueAcrossThreads() throws Exception { + // ValidatesRunner suites build pipelines from many threads at once; a lost update on the + // assertion counter would hand two PAsserts the same name (possibly within one pipeline). + final int threads = 16; + final int namesPerThread = 2000; + final Set names = ConcurrentHashMap.newKeySet(); + final CyclicBarrier start = new CyclicBarrier(threads); + ExecutorService executor = Executors.newFixedThreadPool(threads); + try { + List> futures = new ArrayList<>(); + for (int t = 0; t < threads; t++) { + futures.add( + executor.submit( + () -> { + start.await(); + for (int i = 0; i < namesPerThread; i++) { + assertTrue("duplicate assertion name", names.add(PAssert.nextAssertionName())); + } + return null; + })); + } + for (Future future : futures) { + future.get(); + } + } finally { + executor.shutdownNow(); + } + assertEquals(threads * namesPerThread, names.size()); + } + @Test public void testFailureWithExceptionEncodedDecoded() throws IOException { Throwable error; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestPipelineTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestPipelineTest.java index a79bfd739234..932657ecb218 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestPipelineTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestPipelineTest.java @@ -18,13 +18,18 @@ package org.apache.beam.sdk.testing; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.everyItem; import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.lessThanOrEqualTo; +import static org.hamcrest.Matchers.matchesPattern; import static org.hamcrest.Matchers.not; import static org.hamcrest.Matchers.sameInstance; import static org.hamcrest.Matchers.startsWith; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import com.fasterxml.jackson.databind.ObjectMapper; @@ -44,6 +49,7 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.options.ValueProvider; +import org.apache.beam.sdk.runners.TransformHierarchy; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; @@ -60,8 +66,10 @@ import org.junit.rules.ExpectedException; import org.junit.rules.RuleChain; import org.junit.rules.TestRule; +import org.junit.rules.Timeout; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.junit.runners.model.Statement; /** Tests for {@link TestPipeline}. */ public class TestPipelineTest implements Serializable { @@ -201,6 +209,200 @@ public int hashCode() { } } + /** Tests for {@link TestPipeline#PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES}. */ + @RunWith(JUnit4.class) + public static class UniqueRootNamesTest implements Serializable { + /** Turns unique root names on before the {@link TestPipeline} rule below evaluates. */ + private static final TestRule ENABLE_UNIQUE_ROOT_NAMES = + (base, description) -> + new Statement() { + @Override + public void evaluate() throws Throwable { + System.setProperty( + TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES, "true"); + base.evaluate(); + } + }; + + private final transient TestPipeline pipeline = TestPipeline.create(); + + @Rule + public final transient RuleChain rules = + RuleChain.outerRule(new RestoreSystemProperties()) + .around(ENABLE_UNIQUE_ROOT_NAMES) + .around(pipeline); + + @Test + public void testRuleNamesRootAfterTest() { + String root = pipeline.getRootName(); + assertThat( + root, matchesPattern("t[0-9]+-TestPipelineTest\\$UniqueRootNamesTest-testRuleNames.*")); + assertThat(root.length(), lessThanOrEqualTo(TestPipeline.MAX_ROOT_NAME_LENGTH)); + + PCollection created = pipeline.apply("MyCreate", Create.of(1, 2, 3)); + PCollection mapped = + created.apply( + "MyMap", + MapElements.via( + new SimpleFunction() { + @Override + public Integer apply(Integer input) { + return input; + } + })); + assertThat(created.getName(), startsWith(root + "/")); + assertThat(mapped.getName(), startsWith(root + "/")); + + List fullNames = new ArrayList<>(); + pipeline.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { + if (!node.isRootNode()) { + fullNames.add(node.getFullName()); + } + return CompositeBehavior.ENTER_TRANSFORM; + } + + @Override + public void visitPrimitiveTransform(TransformHierarchy.Node node) { + fullNames.add(node.getFullName()); + } + }); + assertThat(fullNames, hasItem(root + "/MyCreate")); + assertThat(fullNames, hasItem(root + "/MyMap")); + assertThat(fullNames, everyItem(startsWith(root + "/"))); + } + + @Test + public void testRootNamesAreUniqueAcrossPipelines() { + TestPipeline other = TestPipeline.create(); + other.nameRoot("someOtherTest"); + assertThat(other.getRootName(), not(equalTo(pipeline.getRootName()))); + assertThat(other.getRootName(), matchesPattern("t[0-9]+-someOtherTest")); + } + + @Test + public void testNameRootIsIgnoredOnceTransformsExist() { + TestPipeline late = TestPipeline.create(); + late.apply(Create.of(1)); + late.nameRoot("tooLate"); + assertEquals("", late.getRootName()); + } + + @Test + public void testRootNameForSanitizesAndTruncates() { + assertEquals("t7-Foo$Bar-baz", TestPipeline.rootNameFor(7, "Foo$Bar-baz")); + // Characters outside the conservative set, notably the '/' name separator, are replaced. + assertEquals("t7-a_b_c_d", TestPipeline.rootNameFor(7, "a/b c[d")); + // Only the human-readable part is truncated; the unique sequence prefix is kept intact. + String longName = "x".repeat(200); + String truncated = TestPipeline.rootNameFor(123456, longName); + assertEquals(TestPipeline.MAX_ROOT_NAME_LENGTH, truncated.length()); + assertThat(truncated, startsWith("t123456-xxx")); + assertEquals("t1-", TestPipeline.rootNameFor(1, "")); + } + } + + /** Tests for the default (unnamed) root. */ + @RunWith(JUnit4.class) + public static class DefaultRootNameTest { + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + + @Test + public void testRootIsUnnamedUnlessPropertySet() { + assertEquals("", pipeline.getRootName()); + PCollection created = pipeline.apply("MyCreate", Create.of(1)); + assertThat(created.getName(), startsWith("MyCreate")); + } + } + + /** + * Tests for each condition under which {@link TestPipeline#isStandaloneExecutionRequired()} + * reports that the pipeline must not be merged with others. Runs under {@link + * BeamParallelJunit4Runner} because that is what hands the test instance to the rule, which the + * {@link ExpectedException} detection relies on. The {@link Timeout} rule is deliberate: like + * many ValidatesRunner suites this class then executes its test bodies on a thread of JUnit's + * choosing rather than the runner's, and detection must not depend on which thread that is. + */ + @RunWith(BeamParallelJunit4Runner.class) + public static class StandaloneExecutionTest { + @Rule public transient ExpectedException thrown = ExpectedException.none(); + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + @Rule public transient Timeout globalTimeout = Timeout.seconds(60); + + @Test + public void testPlainTestDoesNotRequireStandalone() { + assertFalse(pipeline.isStandaloneExecutionRequired()); + pipeline.apply("MyCreate", Create.of(1)); + assertFalse(pipeline.isStandaloneExecutionRequired()); + } + + @Test(expected = IllegalStateException.class) + public void testExpectedAnnotationRequiresStandalone() { + assertTrue(pipeline.isStandaloneExecutionRequired()); + throw new IllegalStateException("expected"); + } + + @Test + @Category(UsesFailureMessage.class) + public void testUsesFailureMessageCategoryRequiresStandalone() { + assertTrue(pipeline.isStandaloneExecutionRequired()); + } + + @Test + @BeamParallelJunit4Runner.SerialTest + public void testSerialTestMethodRequiresStandalone() { + assertTrue(pipeline.isStandaloneExecutionRequired()); + } + + @Test + public void testExpectedExceptionRuleRequiresStandalone() { + assertFalse(pipeline.isStandaloneExecutionRequired()); + thrown.expect(IllegalStateException.class); + assertTrue(pipeline.isStandaloneExecutionRequired()); + throw new IllegalStateException("expected"); + } + + @Test + public void testOptionsMutatedAfterRuleRequiresStandalone() { + assertFalse(pipeline.isStandaloneExecutionRequired()); + pipeline.getOptions().setJobName("custom-job-name"); + assertTrue(pipeline.isStandaloneExecutionRequired()); + } + + @Test + public void testRuntimeValueProviderRequiresStandalone() { + assertFalse(pipeline.isStandaloneExecutionRequired()); + pipeline.newProvider("runtime-value"); + assertTrue(pipeline.isStandaloneExecutionRequired()); + } + } + + /** Class-level {@link BeamParallelJunit4Runner.SerialTest} also forces standalone execution. */ + @RunWith(JUnit4.class) + @BeamParallelJunit4Runner.SerialTest + public static class SerialClassStandaloneExecutionTest { + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + + @Test + public void testSerialTestClassRequiresStandalone() { + assertTrue(pipeline.isStandaloneExecutionRequired()); + } + } + + /** Class-level {@code @Category(UsesFailureMessage.class)} also forces standalone execution. */ + @RunWith(JUnit4.class) + @Category(UsesFailureMessage.class) + public static class UsesFailureMessageClassStandaloneExecutionTest { + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + + @Test + public void testUsesFailureMessageClassRequiresStandalone() { + assertTrue(pipeline.isStandaloneExecutionRequired()); + } + } + /** * Tests for {@link TestPipeline}'s detection of missing {@link Pipeline#run()}, or abandoned * (dangling) {@link PAssert} or {@link org.apache.beam.sdk.transforms.PTransform} nodes. diff --git a/sdks/java/testing/junit/src/main/java/org/apache/beam/sdk/testing/TestPipelineExtension.java b/sdks/java/testing/junit/src/main/java/org/apache/beam/sdk/testing/TestPipelineExtension.java index 12c05ff47fd1..7bd03ac3d041 100644 --- a/sdks/java/testing/junit/src/main/java/org/apache/beam/sdk/testing/TestPipelineExtension.java +++ b/sdks/java/testing/junit/src/main/java/org/apache/beam/sdk/testing/TestPipelineExtension.java @@ -114,6 +114,9 @@ public void beforeEach(ExtensionContext context) { // Set application name based on test method String appName = getAppName(context); pipeline.getOptions().as(ApplicationNameOptions.class).setAppName(appName); + if (Boolean.getBoolean(TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_UNIQUE_ROOT_NAMES)) { + pipeline.nameRoot(appName); + } // Set up enforcement based on annotations pipeline.setDeducedEnforcementLevel(getAnnotations(context));