gemini-code-assist[bot] commented on code in PR #39341:
URL: https://github.com/apache/beam/pull/39341#discussion_r3588450900
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java:
##########
@@ -45,6 +46,10 @@ public class KafkaStreamsTranslationContext {
private final KafkaStreamsPipelineOptions pipelineOptions;
private final Topology topology;
private final Map<String, String> pCollectionIdToProcessorName;
+ // Accumulates the Beam metrics reported by the SDK harness, one container
per executable stage.
+ // Processors update it as bundles complete (in-JVM reference sharing); the
pipeline result
+ // exposes it as MetricResults.
+ private final MetricsContainerStepMap metricsContainerStepMap = new
MetricsContainerStepMap();
Review Comment:
In a multi-partition Kafka Streams topology, multiple parallel
tasks/partitions will run concurrently. If they all share a single global
`MetricsContainerStepMap` instance, they will concurrently update the same
`MetricsContainerImpl` for a given `transformId`. Since
`CounterCell.update(value)` uses `set(value)` under the hood (overwriting the
value with the latest cumulative bundle metrics from the SDK harness),
concurrent updates from different partitions will overwrite each other instead
of accumulating. We should maintain a thread-safe list of task-specific
`MetricsContainerStepMap`s instead.
```suggestion
// Accumulates the Beam metrics reported by the SDK harness, one container
per executable stage per task.
// Processors update them as bundles complete (in-JVM reference sharing);
the pipeline result
// exposes them as merged MetricResults.
private final List<MetricsContainerStepMap> metricsContainerStepMaps =
new java.util.concurrent.CopyOnWriteArrayList<>();
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java:
##########
@@ -77,6 +82,16 @@ public Topology getTopology() {
return topology;
}
+ /**
+ * Returns the job's metrics accumulator: one {@link
+ * org.apache.beam.runners.core.metrics.MetricsContainerImpl container} per
executable stage,
+ * updated by the stage processors as the SDK harness reports bundle
metrics, and read by the
+ * pipeline result via {@link
MetricsContainerStepMap#asAttemptedOnlyMetricResults}.
+ */
+ public MetricsContainerStepMap getMetricsContainerStepMap() {
+ return metricsContainerStepMap;
+ }
Review Comment:
Expose registration and retrieval methods for the task-specific metrics
accumulators so they can be merged when querying metrics.
```java
/**
* Registers a task-specific metrics accumulator.
*/
public void registerMetricsContainerStepMap(MetricsContainerStepMap map) {
metricsContainerStepMaps.add(map);
}
/**
* Returns the registered task-specific metrics accumulators.
*/
public java.util.List<MetricsContainerStepMap>
getMetricsContainerStepMaps() {
return metricsContainerStepMaps;
}
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java:
##########
@@ -86,12 +86,17 @@ public void translate(
Topology topology = context.getTopology();
// The stage stamps its own transform id on the watermarks it emits, and
aggregates its input
// watermark from the reports of its single upstream transform (the
producer of its input
- // PCollection, whose node name is the upstream transform id).
+ // PCollection, whose node name is the upstream transform id).
Harness-reported metrics land in
+ // this stage's container of the job's metrics step map.
topology.addProcessor(
transformId,
() ->
new ExecutableStageProcessor(
- stagePayload, context.getJobInfo(), transformId,
ImmutableSet.of(parentProcessor)),
+ stagePayload,
+ context.getJobInfo(),
+ transformId,
+ ImmutableSet.of(parentProcessor),
+
context.getMetricsContainerStepMap().getContainer(transformId)),
Review Comment:
Instantiate a new task-specific `MetricsContainerStepMap` inside the
processor supplier (which is called once per task/partition by Kafka Streams)
and register it with the context. This ensures thread safety and prevents
concurrent tasks from overwriting each other's metrics.
```java
() -> {
MetricsContainerStepMap stepMap = new MetricsContainerStepMap();
context.registerMetricsContainerStepMap(stepMap);
return new ExecutableStageProcessor(
stagePayload,
context.getJobInfo(),
transformId,
ImmutableSet.of(parentProcessor),
stepMap.getContainer(transformId));
},
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java:
##########
@@ -62,7 +62,8 @@ public PortablePipelineResult run(RunnerApi.Pipeline
pipeline, JobInfo jobInfo)
KafkaStreams kafkaStreams = new KafkaStreams(topology,
streamsConfig(jobInfo));
kafkaStreams.start();
- return new KafkaStreamsPortablePipelineResult(kafkaStreams);
+ return new KafkaStreamsPortablePipelineResult(
+ kafkaStreams, context.getMetricsContainerStepMap());
Review Comment:
Pass the list of task-specific metrics accumulators to the pipeline result.
```suggestion
return new KafkaStreamsPortablePipelineResult(
kafkaStreams, context.getMetricsContainerStepMaps());
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java:
##########
@@ -21,6 +21,7 @@
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import org.apache.beam.model.jobmanagement.v1.JobApi;
+import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
import org.apache.beam.runners.jobsubmission.PortablePipelineResult;
import org.apache.beam.sdk.metrics.MetricResults;
Review Comment:
Add the necessary imports for merging metrics across multiple task-specific
containers.
```java
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
import org.apache.beam.runners.jobsubmission.PortablePipelineResult;
import org.apache.beam.sdk.metrics.DistributionResult;
import org.apache.beam.sdk.metrics.GaugeResult;
import org.apache.beam.sdk.metrics.MetricKey;
import org.apache.beam.sdk.metrics.MetricName;
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;
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java:
##########
@@ -104,8 +110,9 @@ public State waitUntilFinish() {
@Override
public MetricResults metrics() {
- throw new UnsupportedOperationException(
- "Metrics are not yet implemented in the Kafka Streams runner.");
+ // Attempted values only: the runner does not distinguish committed
results yet (that needs
+ // metrics to be folded into the exactly-once commit, which lands with the
durability work).
+ return
MetricsContainerStepMap.asAttemptedOnlyMetricResults(metricsContainerStepMap);
}
Review Comment:
Implement a thread-safe `MergedMetricResults` that merges counters,
distributions, and gauges across all task-specific accumulators when
`metrics()` is queried.
```java
public MetricResults metrics() {
List<MetricResults> results = new ArrayList<>();
for (MetricsContainerStepMap map : metricsContainerStepMaps) {
results.add(MetricsContainerStepMap.asAttemptedOnlyMetricResults(map));
}
return new MergedMetricResults(results);
}
private static class MergedMetricResults extends MetricResults {
private final List<MetricResults> metricResultsList;
MergedMetricResults(List<MetricResults> metricResultsList) {
this.metricResultsList = metricResultsList;
}
@Override
public MetricQueryResults queryMetrics(MetricsFilter filter) {
List<MetricQueryResults> results = new ArrayList<>();
for (MetricResults mr : metricResultsList) {
results.add(mr.queryMetrics(filter));
}
return mergeQueryResults(results);
}
private MetricQueryResults mergeQueryResults(List<MetricQueryResults>
results) {
Map<MetricKey, MetricResult<Long>> counters = new HashMap<>();
Map<MetricKey, MetricResult<DistributionResult>> distributions = new
HashMap<>();
Map<MetricKey, MetricResult<GaugeResult>> gauges = new HashMap<>();
for (MetricQueryResults qr : results) {
for (MetricResult<Long> counter : qr.getCounters()) {
MetricKey key = MetricKey.create(counter.getStepContext(),
counter.getName());
counters.merge(key, counter, MergedMetricResults::mergeCounters);
}
for (MetricResult<DistributionResult> dist : qr.getDistributions()) {
MetricKey key = MetricKey.create(dist.getStepContext(),
dist.getName());
distributions.merge(key, dist,
MergedMetricResults::mergeDistributions);
}
for (MetricResult<GaugeResult> gauge : qr.getGauges()) {
MetricKey key = MetricKey.create(gauge.getStepContext(),
gauge.getName());
gauges.merge(key, gauge, MergedMetricResults::mergeGauges);
}
}
return new SimpleMetricQueryResults(counters.values(),
distributions.values(), gauges.values());
}
private static MetricResult<Long> mergeCounters(MetricResult<Long> a,
MetricResult<Long> b) {
Long committed = (a.getCommitted() != null || b.getCommitted() != null)
? (a.getCommitted() != null ? a.getCommitted() : 0L) +
(b.getCommitted() != null ? b.getCommitted() : 0L)
: null;
return new SimpleMetricResult<>(
a.getName(),
a.getStepContext(),
committed,
a.getAttempted() + b.getAttempted());
}
private static MetricResult<DistributionResult> mergeDistributions(
MetricResult<DistributionResult> a, MetricResult<DistributionResult>
b) {
DistributionResult committed =
mergeDistributionResults(a.getCommitted(), b.getCommitted());
DistributionResult attempted =
mergeDistributionResults(a.getAttempted(), b.getAttempted());
return new SimpleMetricResult<>(a.getName(), a.getStepContext(),
committed, attempted);
}
private static DistributionResult
mergeDistributionResults(DistributionResult a, DistributionResult b) {
if (a == null) return b;
if (b == null) return a;
return DistributionResult.create(
a.getSum() + b.getSum(),
a.getCount() + b.getCount(),
Math.min(a.getMin(), b.getMin()),
Math.max(a.getMax(), b.getMax()));
}
private static MetricResult<GaugeResult> mergeGauges(
MetricResult<GaugeResult> a, MetricResult<GaugeResult> b) {
GaugeResult committed = mergeGaugeResults(a.getCommitted(),
b.getCommitted());
GaugeResult attempted = mergeGaugeResults(a.getAttempted(),
b.getAttempted());
return new SimpleMetricResult<>(a.getName(), a.getStepContext(),
committed, attempted);
}
private static GaugeResult mergeGaugeResults(GaugeResult a, GaugeResult
b) {
if (a == null) return b;
if (b == null) return a;
if (a.getTimestamp() == null) return b;
if (b.getTimestamp() == null) return a;
return a.getTimestamp().isAfter(b.getTimestamp()) ? a : b;
}
private static class SimpleMetricResult<T> implements MetricResult<T> {
private final MetricName name;
private final String stepContext;
private final T committed;
private final T attempted;
SimpleMetricResult(MetricName name, String stepContext, T committed, T
attempted) {
this.name = name;
this.stepContext = stepContext;
this.committed = committed;
this.attempted = attempted;
}
@Override public MetricName getName() { return name; }
@Override public String getStepContext() { return stepContext; }
@Override public T getCommitted() { return committed; }
@Override public T getAttempted() { return attempted; }
}
private static class SimpleMetricQueryResults implements
MetricQueryResults {
private final Iterable<MetricResult<Long>> counters;
private final Iterable<MetricResult<DistributionResult>> distributions;
private final Iterable<MetricResult<GaugeResult>> gauges;
SimpleMetricQueryResults(
Iterable<MetricResult<Long>> counters,
Iterable<MetricResult<DistributionResult>> distributions,
Iterable<MetricResult<GaugeResult>> gauges) {
this.counters = counters;
this.distributions = distributions;
this.gauges = gauges;
}
@Override public Iterable<MetricResult<Long>> getCounters() { return
counters; }
@Override public Iterable<MetricResult<DistributionResult>>
getDistributions() { return distributions; }
@Override public Iterable<MetricResult<GaugeResult>> getGauges() {
return gauges; }
}
}
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java:
##########
@@ -41,11 +42,16 @@ class KafkaStreamsPortablePipelineResult implements
PortablePipelineResult {
LoggerFactory.getLogger(KafkaStreamsPortablePipelineResult.class);
private final KafkaStreams kafkaStreams;
+ // The job's metrics accumulator, shared by reference with the topology's
stage processors, which
+ // update it as the SDK harness reports bundle metrics.
+ private final MetricsContainerStepMap metricsContainerStepMap;
private final CountDownLatch terminated = new CountDownLatch(1);
private volatile boolean cancelled = false;
- KafkaStreamsPortablePipelineResult(KafkaStreams kafkaStreams) {
+ KafkaStreamsPortablePipelineResult(
+ KafkaStreams kafkaStreams, MetricsContainerStepMap
metricsContainerStepMap) {
this.kafkaStreams = kafkaStreams;
+ this.metricsContainerStepMap = metricsContainerStepMap;
Review Comment:
Update the field and constructor to accept and store the list of
task-specific metrics accumulators.
```suggestion
// The job's metrics accumulators, one per task, updated as the SDK
harness reports bundle metrics.
private final List<MetricsContainerStepMap> metricsContainerStepMaps;
private final CountDownLatch terminated = new CountDownLatch(1);
private volatile boolean cancelled = false;
KafkaStreamsPortablePipelineResult(
KafkaStreams kafkaStreams, List<MetricsContainerStepMap>
metricsContainerStepMaps) {
this.kafkaStreams = kafkaStreams;
this.metricsContainerStepMaps = metricsContainerStepMaps;
```
##########
runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java:
##########
@@ -117,6 +123,8 @@ public static void run(Pipeline pipeline) {
driver.advanceWallClockTime(Duration.ofSeconds(1));
roundTripInternalTopics(driver, internalTopics(topology));
}
+ return MetricsContainerStepMap.asAttemptedOnlyMetricResults(
+ context.getMetricsContainerStepMap());
Review Comment:
Update the test runner to merge metrics across all task-specific
accumulators.
```java
java.util.List<MetricResults> results = new java.util.ArrayList<>();
for (MetricsContainerStepMap map :
context.getMetricsContainerStepMaps()) {
results.add(MetricsContainerStepMap.asAttemptedOnlyMetricResults(map));
}
return new MergedMetricResults(results);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
