Skip to content

Make ValidatesRunner faster: Cache and reuse shared metrics pipeline execution in MetricsTest #40322

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kennknowles merged 1 commit into from
Sep 30, 2026
Merged

Make ValidatesRunner faster: Cache and reuse shared metrics pipeline execution in MetricsTest #40322

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 86 additions & 71 deletions sdks/java/core/src/test/java/org/apache/beam/sdk/metrics/MetricsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import java.io.Serializable;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
Expand Down Expand Up @@ -91,6 +92,9 @@ private static MetricQueryResults queryTestMetrics(PipelineResult result) {

/** Shared test helpers and setup/teardown. */
public abstract static class SharedTestBase implements Serializable {
private static final ConcurrentHashMap<String, PipelineResult> CACHED_METRIC_PIPELINE_RESULTS =
new ConcurrentHashMap<>();

@Rule public final transient ExpectedException thrown = ExpectedException.none();

@Rule public final transient TestPipeline pipeline = TestPipeline.create();
Expand All @@ -101,77 +105,88 @@ public void tearDown() {
}

protected PipelineResult runPipelineWithMetrics() {
final Counter count = Metrics.counter(MetricsTest.class, "count");
StringSet sideinputs = Metrics.stringSet(MetricsTest.class, "sideinputs");
final TupleTag<Integer> output1 = new TupleTag<Integer>() {};
final TupleTag<Integer> output2 = new TupleTag<Integer>() {};
pipeline
.apply(Create.of(5, 8, 13))
.apply(
"MyStep1",
ParDo.of(
new DoFn<Integer, Integer>() {
Distribution bundleDist = Metrics.distribution(MetricsTest.class, "bundle");

@StartBundle
public void startBundle() {
bundleDist.update(10L);
}

@SuppressWarnings("unused")
@ProcessElement
public void processElement(ProcessContext c) {
Distribution values = Metrics.distribution(MetricsTest.class, "input");
StringSet sources = Metrics.stringSet(MetricsTest.class, "sources");
BoundedTrie boundedTrieSources =
Metrics.boundedTrie(MetricsTest.class, "boundedTrieSources");
count.inc();
values.update(c.element());

c.output(c.element());
c.output(c.element());
sources.add("gcs");
sources.add("gcs"); // repeated should appear once
sources.add("gcs", "gcs"); // repeated should appear once
sideinputs.add("bigtable", "spanner");
boundedTrieSources.add(ImmutableList.of("ab_source", "cd_source"));
boundedTrieSources.add(ImmutableList.of("ef_source"));
}

@DoFn.FinishBundle
public void finishBundle() {
bundleDist.update(40L);
}
}))
.apply(
"MyStep2",
ParDo.of(
new DoFn<Integer, Integer>() {
@SuppressWarnings("unused")
@ProcessElement
public void processElement(ProcessContext c) {
Distribution values = Metrics.distribution(MetricsTest.class, "input");
Gauge gauge = Metrics.gauge(MetricsTest.class, "my-gauge");
StringSet sinks = Metrics.stringSet(MetricsTest.class, "sinks");
BoundedTrie boundedTrieSinks =
Metrics.boundedTrie(MetricsTest.class, "boundedTrieSinks");
Integer element = c.element();
count.inc();
values.update(element);
gauge.set(12L);
c.output(element);
sinks.add("bq", "kafka", "kafka"); // repeated should appear once
sideinputs.add("bigtable", "sql");
boundedTrieSinks.add(ImmutableList.of("ab_sink", "cd_sink"));
boundedTrieSinks.add(ImmutableList.of("ef_sink"));
c.output(output2, element);
}
})
.withOutputTags(output1, TupleTagList.of(output2)));
PipelineResult result = pipeline.run();

result.waitUntilFinish();
return result;
String cacheKey =
pipeline.getOptions().getRunner().getName()

@kennknowles kennknowles Sep 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't have a lot of context, but I do have concerns.

  1. There is nothing in the cache key to identify the pipeline. Is there just one pipeline possible here? Reading this file, I guess so? The tests are really poor, with mysterious action-at-a-distance between magic values in the runPipelineWithMetrics helper and the assertGauges and so on. I know you aren't here to fix this...
  2. This is a protected method so presumably used in test suites that extend this one? Frankly this is an anti-pattern that we should have a coding agent eliminate, in favor of simple control flow by calling helper methods. And of course if someone overrides the protected method, they lose the caching and the coherence between the test methods and the run pipeline method.

Honestly, having an agent fully rewrite this file might be nice.

@reuvenlax reuvenlax Sep 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah - I think the bigger problem is that this MetricsTest stuff is old and crufty. Maybe we want to just rewrite it all instead of trying to hack around it.

+ ":"
+ System.getProperty(TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS, "");
synchronized (CACHED_METRIC_PIPELINE_RESULTS) {
PipelineResult cachedResult = CACHED_METRIC_PIPELINE_RESULTS.get(cacheKey);
if (cachedResult != null) {
return cachedResult;
}
final Counter count = Metrics.counter(MetricsTest.class, "count");
StringSet sideinputs = Metrics.stringSet(MetricsTest.class, "sideinputs");
final TupleTag<Integer> output1 = new TupleTag<Integer>() {};
final TupleTag<Integer> output2 = new TupleTag<Integer>() {};
pipeline
.apply(Create.of(5, 8, 13))
.apply(
"MyStep1",
ParDo.of(
new DoFn<Integer, Integer>() {
Distribution bundleDist = Metrics.distribution(MetricsTest.class, "bundle");

@StartBundle
public void startBundle() {
bundleDist.update(10L);
}

@SuppressWarnings("unused")
@ProcessElement
public void processElement(ProcessContext c) {
Distribution values = Metrics.distribution(MetricsTest.class, "input");
StringSet sources = Metrics.stringSet(MetricsTest.class, "sources");
BoundedTrie boundedTrieSources =
Metrics.boundedTrie(MetricsTest.class, "boundedTrieSources");
count.inc();
values.update(c.element());

c.output(c.element());
c.output(c.element());
sources.add("gcs");
sources.add("gcs"); // repeated should appear once
sources.add("gcs", "gcs"); // repeated should appear once
sideinputs.add("bigtable", "spanner");
boundedTrieSources.add(ImmutableList.of("ab_source", "cd_source"));
boundedTrieSources.add(ImmutableList.of("ef_source"));
}

@DoFn.FinishBundle
public void finishBundle() {
bundleDist.update(40L);
}
}))
.apply(
"MyStep2",
ParDo.of(
new DoFn<Integer, Integer>() {
@SuppressWarnings("unused")
@ProcessElement
public void processElement(ProcessContext c) {
Distribution values = Metrics.distribution(MetricsTest.class, "input");
Gauge gauge = Metrics.gauge(MetricsTest.class, "my-gauge");
StringSet sinks = Metrics.stringSet(MetricsTest.class, "sinks");
BoundedTrie boundedTrieSinks =
Metrics.boundedTrie(MetricsTest.class, "boundedTrieSinks");
Integer element = c.element();
count.inc();
values.update(element);
gauge.set(12L);
c.output(element);
sinks.add("bq", "kafka", "kafka"); // repeated should appear once
sideinputs.add("bigtable", "sql");
boundedTrieSinks.add(ImmutableList.of("ab_sink", "cd_sink"));
boundedTrieSinks.add(ImmutableList.of("ef_sink"));
c.output(output2, element);
}
})
.withOutputTags(output1, TupleTagList.of(output2)));
PipelineResult result = pipeline.run();

result.waitUntilFinish();
CACHED_METRIC_PIPELINE_RESULTS.put(cacheKey, result);
return result;
}
}
}

Expand Down
Loading