This is an automated email from the ASF dual-hosted git repository.
scott pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f41b2ef [BEAM-6161] Introduce PCollectionConsumerRegistry and add
ElementCount counters to the java SDK
new fb13cb1 Merge pull request #7272: [BEAM-6161] Introduce
PCollectionConsumerRegistry and add ElementCoun…
f41b2ef is described below
commit f41b2ef319d396479d083046280500ad1cbe0be6
Author: Alex Amato <[email protected]>
AuthorDate: Wed Dec 12 17:37:45 2018 -0800
[BEAM-6161] Introduce PCollectionConsumerRegistry and add ElementCount
counters to the java SDK
---
.../fn-execution/src/main/proto/beam_fn_api.proto | 10 +-
.../beam/runners/core/metrics/LabeledMetrics.java | 35 ++++++
.../runners/core/metrics/MetricsContainerImpl.java | 37 ++++--
.../core/metrics/MonitoringInfoMetricName.java | 124 ++++++++++++++++++++
.../core/metrics/SimpleMonitoringInfoBuilder.java | 33 +++++-
.../core/metrics/SpecMonitoringInfoValidator.java | 1 +
.../runners/core/metrics/LabeledMetricsTest.java | 76 ++++++++++++
.../core/metrics/MetricsContainerImplTest.java | 53 +++++++++
.../core/metrics/MonitoringInfoMetricNameTest.java | 127 +++++++++++++++++++++
.../metrics/SimpleMonitoringInfoBuilderTest.java | 8 +-
.../metrics/SpecMonitoringInfoValidatorTest.java | 2 +-
.../fnexecution/control/RemoteExecutionTest.java | 78 +++++++++++--
.../apache/beam/sdk/metrics/DelegatingCounter.java | 64 +++++++++++
.../org/apache/beam/sdk/metrics/MetricResult.java | 5 +-
.../java/org/apache/beam/sdk/metrics/Metrics.java | 41 -------
.../beam/fn/harness/BeamFnDataReadRunner.java | 25 ++--
.../beam/fn/harness/BeamFnDataWriteRunner.java | 7 +-
.../beam/fn/harness/BoundedSourceRunner.java | 10 +-
.../org/apache/beam/fn/harness/CombineRunners.java | 21 ++--
.../fn/harness/DoFnPTransformRunnerFactory.java | 19 +--
.../org/apache/beam/fn/harness/FlattenRunner.java | 15 +--
.../org/apache/beam/fn/harness/MapFnRunners.java | 20 ++--
.../beam/fn/harness/PTransformRunnerFactory.java | 8 +-
.../fn/harness/control/ProcessBundleHandler.java | 21 ++--
.../harness/data/ElementCountFnDataReceiver.java | 54 +++++++++
.../harness/data/PCollectionConsumerRegistry.java | 103 +++++++++++++++++
.../beam/fn/harness/AssignWindowsRunnerTest.java | 13 +--
.../beam/fn/harness/BeamFnDataReadRunnerTest.java | 13 ++-
.../beam/fn/harness/BeamFnDataWriteRunnerTest.java | 8 +-
.../beam/fn/harness/BoundedSourceRunnerTest.java | 10 +-
.../apache/beam/fn/harness/CombineRunnersTest.java | 31 +++--
.../apache/beam/fn/harness/FlattenRunnerTest.java | 26 ++---
.../beam/fn/harness/FnApiDoFnRunnerTest.java | 106 ++++++++---------
.../apache/beam/fn/harness/MapFnRunnersTest.java | 24 ++--
.../harness/control/ProcessBundleHandlerTest.java | 17 ++-
.../data/ElementCountFnDataReceiverTest.java | 70 ++++++++++++
.../data/PCollectionConsumerRegistryTest.java | 88 ++++++++++++++
37 files changed, 1111 insertions(+), 292 deletions(-)
diff --git a/model/fn-execution/src/main/proto/beam_fn_api.proto
b/model/fn-execution/src/main/proto/beam_fn_api.proto
index 623f23d..378de46 100644
--- a/model/fn-execution/src/main/proto/beam_fn_api.proto
+++ b/model/fn-execution/src/main/proto/beam_fn_api.proto
@@ -322,11 +322,6 @@ message MonitoringInfoSpec {
repeated Annotation annotations = 4;
}
-extend google.protobuf.EnumValueOptions {
- // Enum extension to store the MonitoringInfoSpecs.
- MonitoringInfoSpec monitoring_info_spec = 207174266;
-}
-
// The key name and value string of MonitoringInfo annotations.
message Annotation {
string key = 1;
@@ -348,7 +343,7 @@ message MonitoringInfoSpecs {
ELEMENT_COUNT = 1 [(monitoring_info_spec) = {
urn: "beam:metric:element_count:v1",
type_urn: "beam:metrics:sum_int_64",
- required_labels: [ "PTRANSFORM", "PCOLLECTION" ],
+ required_labels: [ "PCOLLECTION" ],
annotations: [ {
key: "description",
value: "The total elements output to a Pcollection by a PTransform."
@@ -411,6 +406,9 @@ message MonitoringInfoLabelProps {
// specifications, constants, etc.
extend google.protobuf.EnumValueOptions {
MonitoringInfoLabelProps label_props = 127337796; // From: commit 0x7970544.
+
+ // Enum extension to store the MonitoringInfoSpecs.
+ MonitoringInfoSpec monitoring_info_spec = 207174266;
}
message MonitoringInfo {
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/LabeledMetrics.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/LabeledMetrics.java
new file mode 100644
index 0000000..bf6295b
--- /dev/null
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/LabeledMetrics.java
@@ -0,0 +1,35 @@
+/*
+ * 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.core.metrics;
+
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.DelegatingCounter;
+
+/**
+ * Define a metric on the current MetricContainer with a specific URN and a
set of labels. This is a
+ * more convenient way to collect the necessary fields to repackage the metric
into a MonitoringInfo
+ * proto later. Intended for internal use only (SDK and RunnerHarness
developement).
+ */
+public class LabeledMetrics {
+ /**
+ * Create a metric that can be incremented and decremented, and is
aggregated by taking the sum.
+ */
+ public static Counter counter(MonitoringInfoMetricName metricName) {
+ return new DelegatingCounter(metricName);
+ }
+}
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MetricsContainerImpl.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MetricsContainerImpl.java
index bd16121..fa2b834 100644
---
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MetricsContainerImpl.java
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MetricsContainerImpl.java
@@ -22,6 +22,7 @@ import static
org.apache.beam.vendor.guava.v20_0.com.google.common.base.Precondi
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Map;
+import java.util.Map.Entry;
import javax.annotation.Nullable;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfo;
import org.apache.beam.runners.core.construction.metrics.MetricKey;
@@ -138,19 +139,39 @@ public class MetricsContainerImpl implements
Serializable, MetricsContainer {
extractUpdates(counters), extractUpdates(distributions),
extractUpdates(gauges));
}
+ /**
+ * @param metricUpdate
+ * @return The MonitoringInfo generated from the metricUpdate.
+ */
+ private MonitoringInfo counterUpdateToMonitoringInfo(MetricUpdate<Long>
metricUpdate) {
+ SimpleMonitoringInfoBuilder builder = new
SimpleMonitoringInfoBuilder(true);
+ MetricName metricName = metricUpdate.getKey().metricName();
+ if (metricName instanceof MonitoringInfoMetricName) {
+ MonitoringInfoMetricName monitoringInfoName = (MonitoringInfoMetricName)
metricName;
+ // Represents a specific MonitoringInfo for a specific URN.
+ builder.setUrn(monitoringInfoName.getUrn());
+ for (Entry<String, String> e :
monitoringInfoName.getLabels().entrySet()) {
+ builder.setLabel(e.getKey(), e.getValue());
+ }
+ } else {
+ // Represents a user counter.
+ builder.setUrnForUserMetric(
+ metricUpdate.getKey().metricName().getNamespace(),
+ metricUpdate.getKey().metricName().getName());
+ }
+ builder.setInt64Value(metricUpdate.getUpdate());
+ builder.setTimestampToNow();
+ return builder.build();
+ }
+
/** Return the cumulative values for any metrics in this container as
MonitoringInfos. */
public Iterable<MonitoringInfo> getMonitoringInfos() {
// Extract user metrics and store as MonitoringInfos.
ArrayList<MonitoringInfo> monitoringInfos = new
ArrayList<MonitoringInfo>();
- MetricUpdates mus = this.getUpdates();
+ MetricUpdates metricUpdates = this.getUpdates();
- for (MetricUpdate<Long> mu : mus.counterUpdates()) {
- SimpleMonitoringInfoBuilder builder = new
SimpleMonitoringInfoBuilder(true);
- builder.setUrnForUserMetric(
- mu.getKey().metricName().getNamespace(),
mu.getKey().metricName().getName());
- builder.setInt64Value(mu.getUpdate());
- builder.setTimestampToNow();
- monitoringInfos.add(builder.build());
+ for (MetricUpdate<Long> metricUpdate : metricUpdates.counterUpdates()) {
+ monitoringInfos.add(counterUpdateToMonitoringInfo(metricUpdate));
}
return monitoringInfos;
}
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoMetricName.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoMetricName.java
new file mode 100644
index 0000000..5d16910
--- /dev/null
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoMetricName.java
@@ -0,0 +1,124 @@
+/*
+ * 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.core.metrics;
+
+import static
org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map.Entry;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.metrics.MetricName;
+import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Strings;
+
+/**
+ * An implementation of {@code MetricKey} based on a MonitoringInfo's URN and
label to represent the
+ * key instead of only a name+namespace. This is useful when defining system
defined metrics with a
+ * specific urn via a {@code CounterContainer}.
+ */
+public class MonitoringInfoMetricName extends MetricName {
+
+ private String urn;
+ @Nullable private String name;
+ @Nullable private String namespace;
+ private HashMap<String, String> labels = new HashMap<String, String>();
+
+ private MonitoringInfoMetricName(String urn, HashMap<String, String> labels)
{
+ checkArgument(!Strings.isNullOrEmpty(urn), "MonitoringInfoMetricName urn
must be non-empty");
+ checkArgument(labels != null, "MonitoringInfoMetricName labels must be
non-null");
+ // TODO(ajamato): Move SimpleMonitoringInfoBuilder to
beam-runner-core-construction-java
+ // and ensure all necessary labels are set for the specific URN.
+ this.urn = urn;
+ for (Entry<String, String> entry : labels.entrySet()) {
+ this.labels.put(entry.getKey(), entry.getValue());
+ }
+ }
+
+ /** Parse the urn field into a name and namespace field. */
+ private void parseUrn() {
+ if
(this.urn.startsWith(SimpleMonitoringInfoBuilder.USER_COUNTER_URN_PREFIX)) {
+ List<String> split = new
ArrayList<String>(Arrays.asList(this.getUrn().split(":")));
+ this.name = split.get(split.size() - 1);
+ this.namespace = split.get(split.size() - 2);
+ }
+ }
+
+ /** @return the parsed namespace from the user metric URN, otherwise null. */
+ @Override
+ public String getNamespace() {
+ if (this.namespace == null) {
+ parseUrn();
+ }
+ return this.namespace;
+ }
+
+ /** @return the parsed name from the user metric URN, otherwise null. */
+ @Override
+ public String getName() {
+ if (this.name == null) {
+ parseUrn();
+ }
+ return this.name;
+ }
+
+ /** @return the urn of this MonitoringInfo metric. */
+ public String getUrn() {
+ return this.urn;
+ }
+
+ /** @return The labels associated with this MonitoringInfo. */
+ public HashMap<String, String> getLabels() {
+ return this.labels;
+ }
+
+ /** @return a MetricName for a specific urn and labels map. */
+ public static MonitoringInfoMetricName named(String urn, HashMap<String,
String> labels) {
+ return new MonitoringInfoMetricName(urn, labels);
+ }
+
+ @Override
+ public int hashCode() {
+ // Don't include name and namespace, since they are lazily set.
+ return Objects.hash(urn, labels);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ // If the object is compared with itself then return true
+ if (o == this) {
+ return true;
+ }
+ if (!(o instanceof MonitoringInfoMetricName)) {
+ return false;
+ }
+ MonitoringInfoMetricName other = (MonitoringInfoMetricName) o;
+ return this.urn.equals(other.urn) && this.labels.equals(other.labels);
+ }
+
+ @Override
+ public String toString() {
+ StringBuilder builder = new StringBuilder();
+ builder.append(this.urn.toString());
+ builder.append(" ");
+ builder.append(this.labels.toString());
+ return builder.toString();
+ }
+}
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilder.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilder.java
index cdebefc..df1fe50 100644
---
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilder.java
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilder.java
@@ -22,11 +22,14 @@ import java.util.HashMap;
import javax.annotation.Nullable;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfo;
+import
org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfo.MonitoringInfoLabels;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfoLabelProps;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfoSpec;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfoSpecs;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfoTypeUrns;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfoUrns;
import org.apache.beam.runners.core.construction.BeamUrns;
+import
org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -65,6 +68,9 @@ public class SimpleMonitoringInfoBuilder {
private static final HashMap<String, MonitoringInfoSpec> specs =
new HashMap<String, MonitoringInfoSpec>();
+ public static final String PCOLLECTION_LABEL =
getLabelString(MonitoringInfoLabels.PCOLLECTION);
+ public static final String PTRANSFORM_LABEL =
getLabelString(MonitoringInfoLabels.TRANSFORM);
+
private final boolean validateAndDropInvalid;
private static final Logger LOG =
LoggerFactory.getLogger(SimpleMonitoringInfoBuilder.class);
@@ -85,6 +91,13 @@ public class SimpleMonitoringInfoBuilder {
}
}
+ /** Returns the label string constant defined in the MonitoringInfoLabel
enum proto. */
+ private static String getLabelString(MonitoringInfoLabels label) {
+ MonitoringInfoLabelProps props =
+
label.getValueDescriptor().getOptions().getExtension(BeamFnApi.labelProps);
+ return props.getName();
+ }
+
public SimpleMonitoringInfoBuilder() {
this(true);
}
@@ -95,7 +108,7 @@ public class SimpleMonitoringInfoBuilder {
}
/** @return The metric URN for a user metric, with a proper URN prefix. */
- private static String userMetricUrn(String metricNamespace, String
metricName) {
+ public static String userMetricUrn(String metricNamespace, String
metricName) {
String fixedMetricNamespace = metricNamespace.replace(':', '_');
String fixedMetricName = metricName.replace(':', '_');
StringBuilder sb = new StringBuilder();
@@ -144,13 +157,13 @@ public class SimpleMonitoringInfoBuilder {
/** Sets the PTRANSFORM MonitoringInfo label to the given param. */
public SimpleMonitoringInfoBuilder setPTransformLabel(String pTransform) {
// TODO(ajamato): Add validation that it is a valid pTransform name in the
bundle descriptor.
- setLabel("PTRANSFORM", pTransform);
+ setLabel(PTRANSFORM_LABEL, pTransform);
return this;
}
/** Sets the PCOLLECTION MonitoringInfo label to the given param. */
public SimpleMonitoringInfoBuilder setPCollectionLabel(String pCollection) {
- setLabel("PCOLLECTION", pCollection);
+ setLabel(PCOLLECTION_LABEL, pCollection);
return this;
}
@@ -161,17 +174,27 @@ public class SimpleMonitoringInfoBuilder {
}
/**
+ * @return A copy of the MonitoringInfo with the timestamp cleared, to allow
comparing two
+ * MonitoringInfos.
+ */
+ @VisibleForTesting
+ public static MonitoringInfo clearTimestamp(MonitoringInfo input) {
+ MonitoringInfo.Builder builder = MonitoringInfo.newBuilder();
+ builder.mergeFrom(input);
+ builder.clearTimestamp();
+ return builder.build();
+ }
+
+ /**
* Builds the provided MonitoringInfo. Returns null if
validateAndDropInvalid set and fields do
* not match respecting MonitoringInfoSpec based on urn.
*/
@Nullable
public MonitoringInfo build() {
final MonitoringInfo result = this.builder.build();
-
if (validateAndDropInvalid && this.validator.validate(result).isPresent())
{
return null;
}
-
return result;
}
}
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidator.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidator.java
index 07f3acb..b0acbc8 100644
---
a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidator.java
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidator.java
@@ -68,6 +68,7 @@ public class SpecMonitoringInfoValidator {
monitoringInfo.getUrn(), spec.getTypeUrn(),
monitoringInfo.getType()));
}
+ // TODO(ajamato): Tighten this restriction to use set equality, to catch
unused
Set<String> requiredLabels = new HashSet<>(spec.getRequiredLabelsList());
if (!monitoringInfo.getLabelsMap().keySet().containsAll(requiredLabels)) {
return Optional.of(
diff --git
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/LabeledMetricsTest.java
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/LabeledMetricsTest.java
new file mode 100644
index 0000000..c1704aa
--- /dev/null
+++
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/LabeledMetricsTest.java
@@ -0,0 +1,76 @@
+/*
+ * 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.core.metrics;
+
+import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.io.Serializable;
+import java.util.HashMap;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.mockito.Mockito;
+
+/** Tests for {@link LabeledMetrics}. */
+public class LabeledMetricsTest implements Serializable {
+
+ @Rule public final transient ExpectedException thrown =
ExpectedException.none();
+
+ @Test
+ public void testCounterDoesNotFailOperationsWhenNoMetricsContainerPresent() {
+ MetricsEnvironment.setCurrentContainer(null);
+ assertNull(MetricsEnvironment.getCurrentContainer());
+ HashMap<String, String> labels = new HashMap<String, String>();
+ String urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ MonitoringInfoMetricName name = MonitoringInfoMetricName.named(urn,
labels);
+
+ Counter counter = LabeledMetrics.counter(name);
+ counter.inc();
+ counter.inc(5L);
+ counter.dec();
+ counter.dec(5L);
+ }
+
+ @Test
+ public void testOperationsUpdateCounterFromContainerWhenContainerIsPresent()
{
+ HashMap<String, String> labels = new HashMap<String, String>();
+ String urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ MonitoringInfoMetricName name = MonitoringInfoMetricName.named(urn,
labels);
+
+ MetricsContainer mockContainer = Mockito.mock(MetricsContainer.class);
+ Counter mockCounter = Mockito.mock(Counter.class);
+ when(mockContainer.getCounter(name)).thenReturn(mockCounter);
+
+ Counter counter = LabeledMetrics.counter(name);
+
+ MetricsEnvironment.setCurrentContainer(mockContainer);
+ counter.inc();
+ verify(mockCounter).inc(1);
+
+ counter.inc(47L);
+ verify(mockCounter).inc(47);
+
+ counter.dec(5L);
+ verify(mockCounter).inc(-5);
+ }
+}
diff --git
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MetricsContainerImplTest.java
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MetricsContainerImplTest.java
index a4a44d3..ba8749c 100644
---
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MetricsContainerImplTest.java
+++
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MetricsContainerImplTest.java
@@ -24,6 +24,9 @@ import static
org.hamcrest.collection.IsIterableContainingInAnyOrder.containsInA
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
+import java.util.ArrayList;
+import java.util.HashMap;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfo;
import org.apache.beam.sdk.metrics.MetricName;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -144,4 +147,54 @@ public class MetricsContainerImplTest {
DistributionCell dne = container.tryGetDistribution(MetricName.named("ns",
"dne"));
assertEquals(dne, null);
}
+
+ @Test
+ public void testMonitoringInfosArePopulatedForUserCounters() {
+ MetricsContainerImpl testObject = new MetricsContainerImpl("step1");
+ CounterCell c1 = testObject.getCounter(MetricName.named("ns", "name1"));
+ CounterCell c2 = testObject.getCounter(MetricName.named("ns", "name2"));
+ c1.inc(2L);
+ c2.inc(4L);
+ c1.inc(3L);
+
+ SimpleMonitoringInfoBuilder builder1 = new SimpleMonitoringInfoBuilder();
+ builder1.setUrnForUserMetric("ns", "name1");
+ builder1.setInt64Value(5);
+ builder1.build();
+
+ SimpleMonitoringInfoBuilder builder2 = new SimpleMonitoringInfoBuilder();
+ builder2.setUrnForUserMetric("ns", "name2");
+ builder2.setInt64Value(4);
+ builder2.build();
+
+ ArrayList<MonitoringInfo> actualMonitoringInfos = new
ArrayList<MonitoringInfo>();
+ for (MonitoringInfo mi : testObject.getMonitoringInfos()) {
+
actualMonitoringInfos.add(SimpleMonitoringInfoBuilder.clearTimestamp(mi));
+ }
+
+ assertThat(actualMonitoringInfos, containsInAnyOrder(builder1.build(),
builder2.build()));
+ }
+
+ @Test
+ public void testMonitoringInfosArePopulatedForABeamCounter() {
+ MetricsContainerImpl testObject = new MetricsContainerImpl("step1");
+ HashMap<String, String> labels = new HashMap<String, String>();
+ labels.put(SimpleMonitoringInfoBuilder.PCOLLECTION_LABEL, "pcollection");
+ MetricName name =
+
MonitoringInfoMetricName.named(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN,
labels);
+ CounterCell c1 = testObject.getCounter(name);
+ c1.inc(2L);
+
+ SimpleMonitoringInfoBuilder builder1 = new SimpleMonitoringInfoBuilder();
+ builder1.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN);
+ builder1.setPCollectionLabel("pcollection");
+ builder1.setInt64Value(2);
+ builder1.build();
+
+ ArrayList<MonitoringInfo> actualMonitoringInfos = new
ArrayList<MonitoringInfo>();
+ for (MonitoringInfo mi : testObject.getMonitoringInfos()) {
+
actualMonitoringInfos.add(SimpleMonitoringInfoBuilder.clearTimestamp(mi));
+ }
+ assertThat(actualMonitoringInfos, containsInAnyOrder(builder1.build()));
+ }
}
diff --git
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MonitoringInfoMetricNameTest.java
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MonitoringInfoMetricNameTest.java
new file mode 100644
index 0000000..13dd79d
--- /dev/null
+++
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MonitoringInfoMetricNameTest.java
@@ -0,0 +1,127 @@
+/*
+ * 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.core.metrics;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+import java.io.Serializable;
+import java.util.HashMap;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+/** Tests for {@link MonitoringInfoMetricName}. */
+public class MonitoringInfoMetricNameTest implements Serializable {
+
+ @Rule public final transient ExpectedException thrown =
ExpectedException.none();
+
+ @Test
+ public void testElementCountConstruction() {
+ HashMap<String, String> labels = new HashMap<String, String>();
+ String urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ MonitoringInfoMetricName name = MonitoringInfoMetricName.named(urn,
labels);
+ assertEquals(null, name.getName());
+ assertEquals(null, name.getNamespace());
+ assertEquals(labels, name.getLabels());
+ assertEquals(urn, name.getUrn());
+
+ assertEquals(name, name); // test self equals;
+
+ // Reconstruct and test equality and hash code equivalence
+ urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ labels = new HashMap<String, String>();
+ MonitoringInfoMetricName name2 = MonitoringInfoMetricName.named(urn,
labels);
+
+ assertEquals(name, name2);
+ assertEquals(name.hashCode(), name2.hashCode());
+ }
+
+ @Test
+ public void testUserCounterUrnConstruction() {
+ String urn = SimpleMonitoringInfoBuilder.userMetricUrn("namespace",
"name");
+ HashMap<String, String> labels = new HashMap<String, String>();
+ MonitoringInfoMetricName name = MonitoringInfoMetricName.named(urn,
labels);
+ assertEquals("name", name.getName());
+ assertEquals("namespace", name.getNamespace());
+ assertEquals(labels, name.getLabels());
+ assertEquals(urn, name.getUrn());
+
+ assertEquals(name, name); // test self equals;
+
+ // Reconstruct and test equality and hash code equivalence
+ urn = SimpleMonitoringInfoBuilder.userMetricUrn("namespace", "name");
+ labels = new HashMap<String, String>();
+ MonitoringInfoMetricName name2 = MonitoringInfoMetricName.named(urn,
labels);
+
+ assertEquals(name, name2);
+ assertEquals(name.hashCode(), name2.hashCode());
+ }
+
+ @Test
+ public void testNotEqualsDiffLabels() {
+ HashMap<String, String> labels = new HashMap<String, String>();
+ String urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ MonitoringInfoMetricName name = MonitoringInfoMetricName.named(urn,
labels);
+
+ // Reconstruct and test equality and hash code equivalence
+ urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ labels = new HashMap<String, String>();
+ labels.put("label", "value1");
+ MonitoringInfoMetricName name2 = MonitoringInfoMetricName.named(urn,
labels);
+
+ assertNotEquals(name, name2);
+ assertNotEquals(name.hashCode(), name2.hashCode());
+ }
+
+ @Test
+ public void testNotEqualsDiffUrn() {
+ HashMap<String, String> labels = new HashMap<String, String>();
+ String urn = SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN;
+ MonitoringInfoMetricName name = MonitoringInfoMetricName.named(urn,
labels);
+
+ // Reconstruct and test equality and hash code equivalence
+ urn = "differentUrn";
+ labels = new HashMap<String, String>();
+ MonitoringInfoMetricName name2 = MonitoringInfoMetricName.named(urn,
labels);
+
+ assertNotEquals(name, name2);
+ assertNotEquals(name.hashCode(), name2.hashCode());
+ }
+
+ @Test
+ public void testNullLabelsThrows() {
+ thrown.expect(IllegalArgumentException.class);
+ HashMap<String, String> labels = null;
+
MonitoringInfoMetricName.named(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN,
labels);
+ }
+
+ @Test
+ public void testNullUrnThrows() {
+ HashMap<String, String> labels = new HashMap<String, String>();
+ thrown.expect(IllegalArgumentException.class);
+ MonitoringInfoMetricName.named(null, labels);
+ }
+
+ @Test
+ public void testEmptyUrnThrows() {
+ HashMap<String, String> labels = new HashMap<String, String>();
+ thrown.expect(IllegalArgumentException.class);
+ MonitoringInfoMetricName.named("", labels);
+ }
+}
diff --git
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilderTest.java
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilderTest.java
index 2697ef8..14fa703 100644
---
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilderTest.java
+++
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SimpleMonitoringInfoBuilderTest.java
@@ -39,15 +39,13 @@ public class SimpleMonitoringInfoBuilderTest {
builder.setInt64Value(1);
assertNull(builder.build());
- builder.setPTransformLabel("myTransform");
- assertNull(builder.build());
-
builder.setPCollectionLabel("myPcollection");
// Pass now that the spec is fully met.
MonitoringInfo monitoringInfo = builder.build();
assertTrue(monitoringInfo != null);
- assertEquals("myTransform",
monitoringInfo.getLabelsOrDefault("PTRANSFORM", null));
- assertEquals("myPcollection",
monitoringInfo.getLabelsOrDefault("PCOLLECTION", null));
+ assertEquals(
+ "myPcollection",
+
monitoringInfo.getLabelsOrDefault(SimpleMonitoringInfoBuilder.PCOLLECTION_LABEL,
null));
assertEquals(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN,
monitoringInfo.getUrn());
assertEquals(SimpleMonitoringInfoBuilder.SUM_INT64_TYPE_URN,
monitoringInfo.getType());
assertEquals(1,
monitoringInfo.getMetric().getCounterData().getInt64Value());
diff --git
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidatorTest.java
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidatorTest.java
index b764e92..7cff0a0 100644
---
a/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidatorTest.java
+++
b/runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/SpecMonitoringInfoValidatorTest.java
@@ -70,7 +70,7 @@ public class SpecMonitoringInfoValidatorTest {
MonitoringInfo.newBuilder()
.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN)
.setType(SimpleMonitoringInfoBuilder.SUM_INT64_TYPE_URN)
- .putLabels("PCOLLECTION", "anotherValue")
+ .putLabels("PTRANSFORM", "unexpectedLabel")
.build();
assertTrue(testObject.validate(testInput).isPresent());
}
diff --git
a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
index 4c7a56e..09cf59d 100644
---
a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
+++
b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.runners.fnexecution.control;
+import static junit.framework.TestCase.assertEquals;
import static
org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkState;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
@@ -95,6 +96,7 @@ import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.Impulse;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.ParDo.SingleOutput;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
@@ -114,13 +116,11 @@ import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterators;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Sets;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.ThreadFactoryBuilder;
-import org.hamcrest.CoreMatchers;
import org.hamcrest.collection.IsEmptyIterable;
import org.hamcrest.collection.IsIterableContainingInOrder;
import org.joda.time.DateTimeUtils;
import org.joda.time.Duration;
import org.junit.After;
-import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -507,13 +507,36 @@ public class RemoteExecutionTest implements Serializable {
"create",
ParDo.of(
new DoFn<byte[], String>() {
+ private boolean emitted = false;
+
@ProcessElement
public void process(ProcessContext ctxt) {
- Metrics.counter(RemoteExecutionTest.class,
counterMetricName).inc();
+ // TODO(BEAM-6467): Impulse is producing two elements
instead of one.
+ // So add this check to only emit these three
elemenets.
+ if (!emitted) {
+ ctxt.output("zero");
+ ctxt.output("one");
+ ctxt.output("two");
+ Metrics.counter(RemoteExecutionTest.class,
counterMetricName).inc();
+ }
+ emitted = true;
}
}))
.setCoder(StringUtf8Coder.of());
+ SingleOutput<String, String> pardo =
+ ParDo.of(
+ new DoFn<String, String>() {
+ @ProcessElement
+ public void process(ProcessContext ctxt) {
+ // Output the element twice to keep unique numbers in asserts,
6 output elements.
+ ctxt.output(ctxt.element());
+ ctxt.output(ctxt.element());
+ }
+ });
+ input.apply("processA", pardo).setCoder(StringUtf8Coder.of());
+ input.apply("processB", pardo).setCoder(StringUtf8Coder.of());
+
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto);
Optional<ExecutableStage> optionalStage =
@@ -584,21 +607,52 @@ public class RemoteExecutionTest implements Serializable {
@Override
public void onCompleted(ProcessBundleResponse response) {
// Assert the timestamps are non empty then 0 them out before
comparing.
- List<MonitoringInfo> actualMIs = new ArrayList<>();
+ List<MonitoringInfo> result = new ArrayList<MonitoringInfo>();
for (MonitoringInfo mi : response.getMonitoringInfosList()) {
- MonitoringInfo.Builder builder = MonitoringInfo.newBuilder();
- Assert.assertTrue(mi.getTimestamp().getSeconds() > 0);
- builder.mergeFrom(mi);
- builder.clearTimestamp();
- actualMIs.add(builder.build());
+ result.add(SimpleMonitoringInfoBuilder.clearTimestamp(mi));
}
+ // The user counter is counted in both ParDos, so it should be
counted twice for each
+ // element 4 = 2 x 2 elements.
+ List<MonitoringInfo> expected = new ArrayList<MonitoringInfo>();
+
SimpleMonitoringInfoBuilder builder = new
SimpleMonitoringInfoBuilder();
builder.setUrnForUserMetric(RemoteExecutionTest.class.getName(),
counterMetricName);
+ // TODO(ajamato): Add test having a user counter with the same
name, in two
+ // separate ptransforms, should be labelled differnetly. Currently
this does not work
+ // because the MetricContainer is not being seeded properly with
the step name.
+ builder.setInt64Value(1); // Count just the one inc() in create();
+ expected.add(builder.build());
+
+ // The element counter should be counted only once for the
pcollection.
+ // So there should be only two elements.
+ builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN);
+ builder.setPCollectionLabel("impulse.out");
builder.setInt64Value(2);
- MonitoringInfo expectedCounter = builder.build();
-
- assertThat(actualMIs, CoreMatchers.hasItems(expectedCounter));
+ expected.add(builder.build());
+
+ builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN);
+ builder.setPCollectionLabel("create/ParMultiDo(Anonymous).output");
+ builder.setInt64Value(3);
+ expected.add(builder.build());
+
+ // Verify that the element count is not double counted if two
PCollections consume it.
+ builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN);
+
builder.setPCollectionLabel("processA/ParMultiDo(Anonymous).output");
+ builder.setInt64Value(6);
+ expected.add(builder.build());
+
+ builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN);
+
builder.setPCollectionLabel("processB/ParMultiDo(Anonymous).output");
+ builder.setInt64Value(6);
+ expected.add(builder.build());
+
+ assertEquals(5, result.size());
+ assertThat(result, containsInAnyOrder(expected.toArray()));
}
};
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/DelegatingCounter.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/DelegatingCounter.java
new file mode 100644
index 0000000..c2ce493
--- /dev/null
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/DelegatingCounter.java
@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.metrics;
+
+// Intentionally package private, this is meant to be used with Metrics.java
and
+// LabeledMetrics.java
+
+import java.io.Serializable;
+
+/** Implementation of {@link Counter} that delegates to the instance for the
current context. */
+public class DelegatingCounter implements Metric, Counter, Serializable {
+ private final MetricName name;
+
+ public DelegatingCounter(MetricName name) {
+ this.name = name;
+ }
+
+ /** Increment the counter. */
+ @Override
+ public void inc() {
+ inc(1);
+ }
+
+ /** Increment the counter by the given amount. */
+ @Override
+ public void inc(long n) {
+ MetricsContainer container = MetricsEnvironment.getCurrentContainer();
+ if (container != null) {
+ container.getCounter(name).inc(n);
+ }
+ }
+
+ /* Decrement the counter. */
+ @Override
+ public void dec() {
+ inc(-1);
+ }
+
+ /* Decrement the counter by the given amount. */
+ @Override
+ public void dec(long n) {
+ inc(-1 * n);
+ }
+
+ @Override
+ public MetricName getName() {
+ return name;
+ }
+}
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricResult.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricResult.java
index 215528c..d8e0875 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricResult.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricResult.java
@@ -21,7 +21,10 @@ import com.fasterxml.jackson.annotation.JsonFilter;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
-/** The results of a single current metric. */
+/**
+ * The results of a single current metric. TODO(BEAM-6265): Decouple wire
formats from internal
+ * formats, remove usage of MetricName.
+ */
@Experimental(Kind.METRICS)
@JsonFilter("committedMetrics")
public interface MetricResult<T> {
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java
index fd22120..df3666a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java
@@ -98,47 +98,6 @@ public class Metrics {
return new DelegatingGauge(MetricName.named(namespace, name));
}
- /** Implementation of {@link Counter} that delegates to the instance for the
current context. */
- private static class DelegatingCounter implements Metric, Counter,
Serializable {
- private final MetricName name;
-
- private DelegatingCounter(MetricName name) {
- this.name = name;
- }
-
- /** Increment the counter. */
- @Override
- public void inc() {
- inc(1);
- }
-
- /** Increment the counter by the given amount. */
- @Override
- public void inc(long n) {
- MetricsContainer container = MetricsEnvironment.getCurrentContainer();
- if (container != null) {
- container.getCounter(name).inc(n);
- }
- }
-
- /* Decrement the counter. */
- @Override
- public void dec() {
- inc(-1);
- }
-
- /* Decrement the counter by the given amount. */
- @Override
- public void dec(long n) {
- inc(-1 * n);
- }
-
- @Override
- public MetricName getName() {
- return name;
- }
- }
-
/**
* Implementation of {@link Distribution} that delegates to the instance for
the current context.
*/
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
index 9de543e..e44b54b 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
@@ -21,12 +21,11 @@ import static
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Itera
import com.google.auto.service.AutoService;
import java.io.IOException;
-import java.util.Collection;
import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
@@ -46,7 +45,6 @@ import org.apache.beam.sdk.fn.data.RemoteGrpcPortRead;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -85,7 +83,7 @@ public class BeamFnDataReadRunner<OutputT> {
Map<String, PCollection> pCollections,
Map<String, RunnerApi.Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
@@ -107,10 +105,11 @@ public class BeamFnDataReadRunner<OutputT> {
} else {
coderSpec = null;
}
- // TODO make the receiver aware of its transform context as well.
- Collection<FnDataReceiver<WindowedValue<OutputT>>> consumers =
- (Collection)
-
pCollectionIdsToConsumers.get(getOnlyElement(pTransform.getOutputsMap().values()));
+ FnDataReceiver<WindowedValue<OutputT>> consumer =
+ (FnDataReceiver<WindowedValue<OutputT>>)
+ (FnDataReceiver)
+ pCollectionConsumerRegistry.getMultiplexingConsumer(
+ getOnlyElement(pTransform.getOutputsMap().values()));
BeamFnDataReadRunner<OutputT> runner =
new BeamFnDataReadRunner<>(
@@ -120,7 +119,7 @@ public class BeamFnDataReadRunner<OutputT> {
coderSpec,
coders,
beamFnDataClient,
- consumers);
+ consumer);
startFunctionRegistry.register(pTransformId,
runner::registerInputLocation);
finishFunctionRegistry.register(pTransformId,
runner::blockTillReadFinishes);
return runner;
@@ -128,7 +127,7 @@ public class BeamFnDataReadRunner<OutputT> {
}
private final Endpoints.ApiServiceDescriptor apiServiceDescriptor;
- private final FnDataReceiver<WindowedValue<OutputT>> receiver;
+ private final FnDataReceiver<WindowedValue<OutputT>> consumer;
private final Supplier<String> processBundleInstructionIdSupplier;
private final BeamFnDataClient beamFnDataClient;
private final Coder<WindowedValue<OutputT>> coder;
@@ -143,14 +142,14 @@ public class BeamFnDataReadRunner<OutputT> {
RunnerApi.Coder coderSpec,
Map<String, RunnerApi.Coder> coders,
BeamFnDataClient beamFnDataClient,
- Collection<FnDataReceiver<WindowedValue<OutputT>>> consumers)
+ FnDataReceiver<WindowedValue<OutputT>> consumer)
throws IOException {
RemoteGrpcPort port =
RemoteGrpcPortRead.fromPTransform(grpcReadNode).getPort();
this.apiServiceDescriptor = port.getApiServiceDescriptor();
this.inputTarget = inputTarget;
this.processBundleInstructionIdSupplier =
processBundleInstructionIdSupplier;
this.beamFnDataClient = beamFnDataClient;
- this.receiver = MultiplexingFnDataReceiver.forConsumers(consumers);
+ this.consumer = consumer;
RehydratedComponents components =
RehydratedComponents.forComponents(Components.newBuilder().putAllCoders(coders).build());
@@ -173,7 +172,7 @@ public class BeamFnDataReadRunner<OutputT> {
apiServiceDescriptor,
LogicalEndpoint.of(processBundleInstructionIdSupplier.get(),
inputTarget),
coder,
- receiver);
+ consumer);
}
public void blockTillReadFinishes() throws Exception {
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
index fae1653..4177814 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
@@ -25,6 +25,7 @@ import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
@@ -44,7 +45,6 @@ import org.apache.beam.sdk.fn.data.RemoteGrpcPortWrite;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -83,7 +83,7 @@ public class BeamFnDataWriteRunner<InputT> {
Map<String, PCollection> pCollections,
Map<String, RunnerApi.Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
@@ -108,9 +108,10 @@ public class BeamFnDataWriteRunner<InputT> {
new BeamFnDataWriteRunner<>(
pTransform, processBundleInstructionId, target, coderSpec,
coders, beamFnDataClient);
startFunctionRegistry.register(pTransformId, runner::registerForOutput);
- pCollectionIdsToConsumers.put(
+ pCollectionConsumerRegistry.register(
getOnlyElement(pTransform.getInputsMap().values()),
(FnDataReceiver) (FnDataReceiver<WindowedValue<InputT>>)
runner::consume);
+
finishFunctionRegistry.register(pTransformId, runner::close);
return runner;
}
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
index ac37582..4923afd 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
@@ -25,6 +25,7 @@ import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.control.ProcessBundleHandler;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -43,7 +44,6 @@ import org.apache.beam.sdk.util.WindowedValue;
import
org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.InvalidProtocolBufferException;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
/**
* A runner which creates {@link Reader}s for each {@link BoundedSource} sent
as an input and
@@ -77,15 +77,13 @@ public class BoundedSourceRunner<InputT extends
BoundedSource<OutputT>, OutputT>
Map<String, PCollection> pCollections,
Map<String, Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener) {
-
- // TODO make the receiver aware of its transform context as well.
ImmutableList.Builder<FnDataReceiver<WindowedValue<?>>> consumers =
ImmutableList.builder();
for (String pCollectionId : pTransform.getOutputsMap().values()) {
- consumers.addAll(pCollectionIdsToConsumers.get(pCollectionId));
+
consumers.add(pCollectionConsumerRegistry.getMultiplexingConsumer(pCollectionId));
}
@SuppressWarnings({"rawtypes", "unchecked"})
@@ -97,7 +95,7 @@ public class BoundedSourceRunner<InputT extends
BoundedSource<OutputT>, OutputT>
FnDataReceiver runReadLoop = (FnDataReceiver<WindowedValue<InputT>>)
runner::runReadLoop;
for (String pCollectionId : pTransform.getInputsMap().values()) {
- pCollectionIdsToConsumers.put(pCollectionId, runReadLoop);
+ pCollectionConsumerRegistry.register(pCollectionId, runReadLoop);
}
return runner;
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
index 9a3d936..a6c4e50 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
@@ -19,12 +19,11 @@ package org.apache.beam.fn.harness;
import com.google.auto.service.AutoService;
import java.io.IOException;
-import java.util.Collection;
import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -47,7 +46,6 @@ import org.apache.beam.sdk.values.KV;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
/** Executes different components of Combine PTransforms. */
public class CombineRunners {
@@ -124,7 +122,7 @@ public class CombineRunners {
Map<String, PCollection> pCollections,
Map<String, RunnerApi.Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
@@ -160,24 +158,19 @@ public class CombineRunners {
Coder<AccumT> accumCoder =
(Coder<AccumT>)
rehydratedComponents.getCoder(combinePayload.getAccumulatorCoderId());
- Collection<FnDataReceiver<WindowedValue<KV<KeyT, AccumT>>>> consumers =
- (Collection)
- pCollectionIdsToConsumers.get(
+ FnDataReceiver<WindowedValue<KV<KeyT, AccumT>>> consumer =
+ (FnDataReceiver)
+ pCollectionConsumerRegistry.getMultiplexingConsumer(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()));
// TODO make the receiver aware of its transform context as well.
// Create the runner.
PrecombineRunner<KeyT, InputT, AccumT> runner =
- new PrecombineRunner<>(
- pipelineOptions,
- combineFn,
- MultiplexingFnDataReceiver.forConsumers(consumers),
- keyCoder,
- accumCoder);
+ new PrecombineRunner<>(pipelineOptions, combineFn, consumer,
keyCoder, accumCoder);
// Register the appropriate handlers.
startFunctionRegistry.register(pTransformId, runner::startBundle);
- pCollectionIdsToConsumers.put(
+ pCollectionConsumerRegistry.register(
Iterables.getOnlyElement(pTransform.getInputsMap().values()),
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<KeyT, InputT>>>)
runner::processElement);
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java
index 7bf0f2c..f4d478d 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java
@@ -24,6 +24,7 @@ import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.fn.harness.state.SideInputSpec;
@@ -88,7 +89,7 @@ abstract class DoFnPTransformRunnerFactory<
Map<String, PCollection> pCollections,
Map<String, RunnerApi.Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener) {
@@ -102,7 +103,7 @@ abstract class DoFnPTransformRunnerFactory<
pCollections,
coders,
windowingStrategies,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
splitListener);
RunnerT runner = createRunner(context);
@@ -116,7 +117,7 @@ abstract class DoFnPTransformRunnerFactory<
context.parDoPayload.getSideInputsMap().keySet(),
context.parDoPayload.getTimerSpecsMap().keySet()));
for (String localInputName : mainInput) {
- pCollectionIdsToConsumers.put(
+ pCollectionConsumerRegistry.register(
pTransform.getInputsOrThrow(localInputName),
(FnDataReceiver) (FnDataReceiver<WindowedValue<TransformInputT>>)
runner::processElement);
}
@@ -127,10 +128,12 @@ abstract class DoFnPTransformRunnerFactory<
DoFnSignatures.getTimerSpecOrThrow(
context.doFnSignature.timerDeclarations().get(localName),
context.doFn)
.getTimeDomain();
- pCollectionIdsToConsumers.put(
+ pCollectionConsumerRegistry.register(
pTransform.getInputsOrThrow(localName),
- timer ->
- runner.processTimer(localName, timeDomain,
(WindowedValue<KV<Object, Timer>>) timer));
+ (FnDataReceiver)
+ timer ->
+ runner.processTimer(
+ localName, timeDomain, (WindowedValue<KV<Object,
Timer>>) timer));
}
finishFunctionRegistry.register(pTransformId, runner::finishBundle);
@@ -170,7 +173,7 @@ abstract class DoFnPTransformRunnerFactory<
Map<String, PCollection> pCollections,
Map<String, RunnerApi.Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
BundleSplitListener splitListener) {
this.pipelineOptions = pipelineOptions;
this.beamFnStateClient = beamFnStateClient;
@@ -279,7 +282,7 @@ abstract class DoFnPTransformRunnerFactory<
localNameToConsumerBuilder = ImmutableListMultimap.builder();
for (Map.Entry<String, String> entry :
pTransform.getOutputsMap().entrySet()) {
localNameToConsumerBuilder.putAll(
- entry.getKey(), pCollectionIdsToConsumers.get(entry.getValue()));
+ entry.getKey(),
pCollectionConsumerRegistry.getMultiplexingConsumer(entry.getValue()));
}
localNameToConsumer = localNameToConsumerBuilder.build();
tagToSideInputSpecMap = tagToSideInputSpecMapBuilder.build();
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FlattenRunner.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FlattenRunner.java
index 6a06780..e75b92f 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FlattenRunner.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FlattenRunner.java
@@ -25,7 +25,7 @@ import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -36,8 +36,6 @@ import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSet;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
/** Executes flatten PTransforms. */
public class FlattenRunner<InputT> {
@@ -64,25 +62,22 @@ public class FlattenRunner<InputT> {
Map<String, PCollection> pCollections,
Map<String, Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
throws IOException {
// Give each input a MultiplexingFnDataReceiver to all outputs of the
flatten.
- ImmutableSet.Builder<FnDataReceiver<WindowedValue<InputT>>>
consumersBuilder =
- new ImmutableSet.Builder<>();
String output = getOnlyElement(pTransform.getOutputsMap().values());
- consumersBuilder.addAll((Iterable)
pCollectionIdsToConsumers.get(output));
+ FnDataReceiver<WindowedValue<?>> receiver =
+ pCollectionConsumerRegistry.getMultiplexingConsumer(output);
- FnDataReceiver<WindowedValue<InputT>> receiver =
- MultiplexingFnDataReceiver.forConsumers(consumersBuilder.build());
FlattenRunner<InputT> runner = new FlattenRunner<>();
// TODO make the receiver aware of its transform context as well.
for (String pCollectionId : pTransform.getInputsMap().values()) {
- pCollectionIdsToConsumers.put(pCollectionId, (FnDataReceiver)
receiver);
+ pCollectionConsumerRegistry.register(pCollectionId, (FnDataReceiver)
receiver);
}
return runner;
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/MapFnRunners.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/MapFnRunners.java
index f3afcd9..6f21506 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/MapFnRunners.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/MapFnRunners.java
@@ -20,12 +20,11 @@ package org.apache.beam.fn.harness;
import static
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables.getOnlyElement;
import java.io.IOException;
-import java.util.Collection;
import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -36,7 +35,6 @@ import org.apache.beam.sdk.fn.function.ThrowingFunction;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
/**
* Utilities to create {@code PTransformRunners} which execute simple map
functions.
@@ -106,22 +104,20 @@ public abstract class MapFnRunners {
Map<String, PCollection> pCollections,
Map<String, RunnerApi.Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
throws IOException {
- Collection<FnDataReceiver<WindowedValue<OutputT>>> consumers =
- (Collection)
-
pCollectionIdsToConsumers.get(getOnlyElement(pTransform.getOutputsMap().values()));
+ FnDataReceiver<WindowedValue<InputT>> consumer =
+ (FnDataReceiver)
+ pCollectionConsumerRegistry.getMultiplexingConsumer(
+ getOnlyElement(pTransform.getOutputsMap().values()));
- // TODO make the receiver aware of its transform context as well.
- Mapper<InputT, OutputT> mapper =
- mapperFactory.create(
- pTransformId, pTransform,
MultiplexingFnDataReceiver.forConsumers(consumers));
+ Mapper<InputT, OutputT> mapper = mapperFactory.create(pTransformId,
pTransform, consumer);
- pCollectionIdsToConsumers.put(
+ pCollectionConsumerRegistry.register(
Iterables.getOnlyElement(pTransform.getInputsMap().values()),
(FnDataReceiver) (FnDataReceiver<WindowedValue<InputT>>)
mapper::map);
return mapper;
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
index a7f129e..ecd0f64 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
@@ -22,16 +22,14 @@ import java.util.Map;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Coder;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.util.WindowedValue;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
/** A factory able to instantiate an appropriate handler for a given
PTransform. */
public interface PTransformRunnerFactory<T> {
@@ -50,7 +48,7 @@ public interface PTransformRunnerFactory<T> {
* @param pCollections A mapping from PCollection id to PCollection
definition.
* @param coders A mapping from coder id to coder definition.
* @param windowingStrategies
- * @param pCollectionIdsToConsumers A mapping from PCollection id to a
collection of consumers.
+ * @param pCollectionConsumerRegistry A mapping from PCollection id to a
collection of consumers.
* Note that if this handler is a consumer, it should register itself
within this multimap
* under the appropriate PCollection ids. Also note that all output
consumers needed by this
* PTransform (based on the values of the {@link
PTransform#getOutputsMap()} will have already
@@ -69,7 +67,7 @@ public interface PTransformRunnerFactory<T> {
Map<String, PCollection> pCollections,
Map<String, Coder> coders,
Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
index 0f3988e..26f30a8 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
@@ -31,6 +31,7 @@ import java.util.function.Supplier;
import org.apache.beam.fn.harness.PTransformRunnerFactory;
import org.apache.beam.fn.harness.PTransformRunnerFactory.Registrar;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.data.QueueingBeamFnDataClient;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
@@ -53,11 +54,9 @@ import
org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowingStrategy;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.fn.function.ThrowingRunnable;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.common.ReflectHelpers;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.Message;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.TextFormat;
@@ -65,7 +64,6 @@ import
org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleF
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.HashMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Lists;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Multimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.SetMultimap;
@@ -148,7 +146,7 @@ public class ProcessBundleHandler {
Supplier<String> processBundleInstructionId,
ProcessBundleDescriptor processBundleDescriptor,
SetMultimap<String, String> pCollectionIdsToConsumingPTransforms,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
Set<String> processedPTransformIds,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
@@ -169,7 +167,7 @@ public class ProcessBundleHandler {
processBundleInstructionId,
processBundleDescriptor,
pCollectionIdsToConsumingPTransforms,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
processedPTransformIds,
startFunctionRegistry,
finishFunctionRegistry,
@@ -202,7 +200,7 @@ public class ProcessBundleHandler {
processBundleDescriptor.getPcollectionsMap(),
processBundleDescriptor.getCodersMap(),
processBundleDescriptor.getWindowingStrategiesMap(),
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
startFunctionRegistry,
finishFunctionRegistry,
splitListener);
@@ -226,8 +224,7 @@ public class ProcessBundleHandler {
(BeamFnApi.ProcessBundleDescriptor) fnApiRegistry.apply(bundleId);
SetMultimap<String, String> pCollectionIdsToConsumingPTransforms =
HashMultimap.create();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers =
- ArrayListMultimap.create();
+ PCollectionConsumerRegistry pCollectionConsumerRegistry = new
PCollectionConsumerRegistry();
HashSet<String> processedPTransformIds = new HashSet<>();
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
@@ -289,13 +286,17 @@ public class ProcessBundleHandler {
request::getInstructionId,
bundleDescriptor,
pCollectionIdsToConsumingPTransforms,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
processedPTransformIds,
startFunctionRegistry,
finishFunctionRegistry,
splitListener);
}
+ // TODO(ajamato): A MetricsContainerImpl should be created and used for
each PTransform
+ // or the metric should be registered in some other way with the
ptransform name context,
+ // not for the instruction. This fails to associate each user counter
with an appropriate
+ // ptransform.
MetricsContainerImpl metricsContainer = new
MetricsContainerImpl(request.getInstructionId());
try (Closeable closeable =
MetricsEnvironment.scopedMetricsContainer(metricsContainer)) {
@@ -413,7 +414,7 @@ public class ProcessBundleHandler {
Map<String, PCollection> pCollections,
Map<String, Coder> coders,
Map<String, WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener) {
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiver.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiver.java
new file mode 100644
index 0000000..5de116e
--- /dev/null
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiver.java
@@ -0,0 +1,54 @@
+/*
+ * 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.fn.harness.data;
+
+import java.util.HashMap;
+import org.apache.beam.runners.core.metrics.LabeledMetrics;
+import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
+import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.util.WindowedValue;
+
+/**
+ * A wrapping {@code FnDataReceiver<WindowedValue<T>>} which counts the number
of elements consumed
+ * by the original {@code FnDataReceiver<WindowedValue<T>>}.
+ *
+ * @param <T> - The receiving type of the PTransform.
+ */
+public class ElementCountFnDataReceiver<T> implements
FnDataReceiver<WindowedValue<T>> {
+
+ private FnDataReceiver<WindowedValue<T>> original;
+ private Counter counter;
+
+ public ElementCountFnDataReceiver(FnDataReceiver<WindowedValue<T>> original,
String pCollection) {
+ this.original = original;
+ HashMap<String, String> labels = new HashMap<String, String>();
+ labels.put(SimpleMonitoringInfoBuilder.PCOLLECTION_LABEL, pCollection);
+ MonitoringInfoMetricName metricName =
+
MonitoringInfoMetricName.named(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN,
labels);
+ this.counter = LabeledMetrics.counter(metricName);
+ }
+
+ @Override
+ public void accept(WindowedValue<T> input) throws Exception {
+ // Increment the counter for each window the element occurs in.
+ this.counter.inc(input.getWindows().size());
+ this.original.accept(input);
+ }
+}
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java
new file mode 100644
index 0000000..8b8dbc4
--- /dev/null
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java
@@ -0,0 +1,103 @@
+/*
+ * 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.fn.harness.data;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.util.WindowedValue;
+import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
+import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
+
+/**
+ * The {@code PCollectionConsumerRegistry} is used to maintain a collection of
consuming
+ * FnDataReceiver for each pCollectionId. Registering with this class allows
inserting an element
+ * count counter for every pCollection. A combined MultiplexingConsumer
(Wrapped with an
+ * ElementCountFnDataReceiver) is returned by calling getMultiplexingConsumer.
+ */
+public class PCollectionConsumerRegistry {
+
+ private ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
pCollectionIdsToConsumers;
+ private Map<String, ElementCountFnDataReceiver>
pCollectionIdsToWrappedConsumer;
+
+ public PCollectionConsumerRegistry() {
+ pCollectionIdsToConsumers = ArrayListMultimap.create();
+ pCollectionIdsToWrappedConsumer = new HashMap<String,
ElementCountFnDataReceiver>();
+ }
+
+ /**
+ * Register the specified consumer to handle the elements in the pCollection
associated with
+ * pCollectionId. All consumers must be registered before extracting the
combined consumer by
+ * calling getMultiplexingConsumer(), or an exception will be thrown.
+ *
+ * @param pCollectionId
+ * @param consumer
+ * @param <T> the element type of the PCollection
+ * @throws RuntimeException if {@code register()} is called after {@code
+ * getMultiplexingConsumer()} is called.
+ */
+ public <T> void register(String pCollectionId,
FnDataReceiver<WindowedValue<T>> consumer) {
+ // Just save these consumers for now, but package them up later with an
+ // ElementCountFnDataReceiver and possibly a MultiplexingFnDataReceiver
+ // if there are multiple consumers.
+ ElementCountFnDataReceiver wrappedConsumer =
+ pCollectionIdsToWrappedConsumer.getOrDefault(pCollectionId, null);
+ if (wrappedConsumer != null) {
+ throw new RuntimeException(
+ "New consumers for a pCollectionId cannot be register()-d after "
+ + "calling getMultiplexingConsumer.");
+ }
+ pCollectionIdsToConsumers.put(pCollectionId, (FnDataReceiver) consumer);
+ }
+
+ /** @return the list of pcollection ids. */
+ public Set<String> keySet() {
+ return pCollectionIdsToConsumers.keySet();
+ }
+
+ /**
+ * New consumers should not be register()-ed after calling this method. This
will cause a
+ * RuntimeException, as this would fail to properly wrap the late-added
consumer to the
+ * ElementCountFnDataReceiver.
+ *
+ * @return A single ElementCountFnDataReceiver which directly wraps all the
registered consumers.
+ */
+ public FnDataReceiver<WindowedValue<?>> getMultiplexingConsumer(String
pCollectionId) {
+ ElementCountFnDataReceiver wrappedConsumer =
+ pCollectionIdsToWrappedConsumer.getOrDefault(pCollectionId, null);
+ if (wrappedConsumer == null) {
+ List<FnDataReceiver<WindowedValue<?>>> consumers =
+ pCollectionIdsToConsumers.get(pCollectionId);
+ FnDataReceiver<WindowedValue<?>> consumer =
+ MultiplexingFnDataReceiver.forConsumers(consumers);
+ wrappedConsumer = new ElementCountFnDataReceiver(consumer,
pCollectionId);
+ pCollectionIdsToWrappedConsumer.put(pCollectionId, wrappedConsumer);
+ }
+ return wrappedConsumer;
+ }
+
+ /**
+ * @return the number of underlying consumers for a pCollectionId, some
tests may wish to check
+ * this.
+ */
+ public List<FnDataReceiver<WindowedValue<?>>> getUnderlyingConsumers(String
pCollectionId) {
+ return pCollectionIdsToConsumers.get(pCollectionId);
+ }
+}
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
index e9f3315..ce0a732 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
@@ -26,6 +26,7 @@ import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import
org.apache.beam.fn.harness.AssignWindowsRunner.AssignWindowsMapFnFactory;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowIntoPayload;
@@ -34,7 +35,6 @@ import
org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.SdkComponents;
import org.apache.beam.runners.core.construction.WindowingStrategyTranslation;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.fn.function.ThrowingFunction;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
@@ -48,11 +48,8 @@ import
org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
import org.apache.beam.sdk.util.WindowedValue;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSet;
-import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Rule;
@@ -172,8 +169,8 @@ public class AssignWindowsRunnerTest implements
Serializable {
}
};
Collection<WindowedValue<?>> outputs = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> receivers =
ArrayListMultimap.create();
- receivers.put("output", outputs::add);
+ PCollectionConsumerRegistry pCollectionConsumerRegistry = new
PCollectionConsumerRegistry();
+ pCollectionConsumerRegistry.register("output", outputs::add);
SdkComponents components = SdkComponents.create();
components.registerEnvironment(Environments.createDockerEnvironment("java"));
MapFnRunners.forWindowedValueMapFnFactory(new
AssignWindowsMapFnFactory<>())
@@ -199,7 +196,7 @@ public class AssignWindowsRunnerTest implements
Serializable {
null /* pCollections */,
null /* coders */,
null /* windowingStrategies */,
- receivers,
+ pCollectionConsumerRegistry,
null /* startFunctionRegistry */,
null, /* finishFunctionRegistry */
null /* splitListener */);
@@ -212,7 +209,7 @@ public class AssignWindowsRunnerTest implements
Serializable {
new IntervalWindow(new Instant(-22L),
Duration.standardMinutes(5L)),
new IntervalWindow(new Instant(-120000L),
Duration.standardMinutes(3L))),
PaneInfo.ON_TIME_AND_ONLY_FIRING);
- Iterables.getOnlyElement(receivers.get("input")).accept(value);
+ pCollectionConsumerRegistry.getMultiplexingConsumer("input").accept(value);
assertThat(
outputs,
containsInAnyOrder(
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
index f4e34a3..d6d2769 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
@@ -39,6 +39,8 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.beam.fn.harness.PTransformRunnerFactory.Registrar;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.pipeline.v1.Endpoints;
@@ -58,11 +60,9 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Suppliers;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.Uninterruptibles;
import org.hamcrest.collection.IsMapContaining;
import org.junit.Before;
@@ -126,9 +126,9 @@ public class BeamFnDataReadRunnerTest {
List<WindowedValue<String>> outputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
String localOutputId = "outputPC";
- consumers.put(
+ consumers.register(
localOutputId, (FnDataReceiver)
(FnDataReceiver<WindowedValue<String>>) outputValues::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry = new
PTransformFunctionRegistry();
@@ -193,7 +193,8 @@ public class BeamFnDataReadRunnerTest {
.thenReturn(bundle2Future);
List<WindowedValue<String>> valuesA = new ArrayList<>();
List<WindowedValue<String>> valuesB = new ArrayList<>();
-
+ FnDataReceiver<WindowedValue<String>> consumers =
+ MultiplexingFnDataReceiver.forConsumers(ImmutableList.of(valuesA::add,
valuesB::add));
AtomicReference<String> bundleId = new AtomicReference<>("0");
BeamFnDataReadRunner<String> readRunner =
new BeamFnDataReadRunner<>(
@@ -203,7 +204,7 @@ public class BeamFnDataReadRunnerTest {
CODER_SPEC,
COMPONENTS.getCodersMap(),
mockBeamFnDataClient,
- ImmutableList.of(valuesA::add, valuesB::add));
+ consumers);
// Process for bundle id 0
readRunner.registerInputLocation();
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
index bc012a1..3f8ed0d 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
@@ -39,6 +39,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.beam.fn.harness.PTransformRunnerFactory.Registrar;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.pipeline.v1.Endpoints;
@@ -48,17 +49,14 @@ import
org.apache.beam.runners.core.construction.CoderTranslation;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.fn.data.LogicalEndpoint;
import org.apache.beam.sdk.fn.data.RemoteGrpcPortWrite;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Suppliers;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.hamcrest.collection.IsMapContaining;
import org.junit.Before;
import org.junit.Test;
@@ -116,7 +114,7 @@ public class BeamFnDataWriteRunnerTest {
public void testCreatingAndProcessingBeamFnDataWriteRunner() throws
Exception {
String bundleId = "57L";
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry = new
PTransformFunctionRegistry();
@@ -181,7 +179,7 @@ public class BeamFnDataWriteRunnerTest {
eq(WIRE_CODER));
assertThat(consumers.keySet(), containsInAnyOrder(localInputId));
-
Iterables.getOnlyElement(consumers.get(localInputId)).accept(valueInGlobalWindow("TestValue"));
+
consumers.getMultiplexingConsumer(localInputId).accept(valueInGlobalWindow("TestValue"));
assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
outputValues.clear();
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
index 240507d..5214b24 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
@@ -31,6 +31,7 @@ import java.util.List;
import java.util.ServiceLoader;
import org.apache.beam.fn.harness.PTransformRunnerFactory.Registrar;
import org.apache.beam.fn.harness.control.ProcessBundleHandler;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
@@ -41,10 +42,8 @@ import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Suppliers;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.hamcrest.Matchers;
import org.hamcrest.collection.IsMapContaining;
import org.junit.Test;
@@ -125,8 +124,8 @@ public class BoundedSourceRunnerTest {
public void testCreatingAndProcessingSourceFromFactory() throws Exception {
List<WindowedValue<String>> outputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
"outputPC", (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
outputValues::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry = new
PTransformFunctionRegistry();
@@ -171,7 +170,8 @@ public class BoundedSourceRunnerTest {
// Check that when passing a source along as an input, the source is
processed.
assertThat(consumers.keySet(), containsInAnyOrder("inputPC", "outputPC"));
- Iterables.getOnlyElement(consumers.get("inputPC"))
+ consumers
+ .getMultiplexingConsumer("inputPC")
.accept(valueInGlobalWindow(CountingSource.upTo(2)));
assertThat(outputValues, contains(valueInGlobalWindow(0L),
valueInGlobalWindow(1L)));
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
index 05cc948..21f6c4d 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
@@ -28,6 +28,7 @@ import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.Collections;
import java.util.Deque;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.construction.PipelineTranslation;
@@ -44,9 +45,7 @@ import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -121,9 +120,9 @@ public class CombineRunnersTest {
@Test
public void testPrecombine() throws Exception {
// Create a map of consumers and an output target to check output values.
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new
ArrayDeque<>();
- consumers.put(
+ consumers.register(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()),
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<String, Integer>>>)
mainOutputValues::add);
@@ -154,8 +153,7 @@ public class CombineRunnersTest {
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
- FnDataReceiver<WindowedValue<?>> input =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ FnDataReceiver<WindowedValue<?>> input =
consumers.getMultiplexingConsumer(inputPCollectionId);
input.accept(valueInGlobalWindow(KV.of("A", "1")));
input.accept(valueInGlobalWindow(KV.of("A", "2")));
input.accept(valueInGlobalWindow(KV.of("A", "6")));
@@ -188,9 +186,9 @@ public class CombineRunnersTest {
@Test
public void testMergeAccumulators() throws Exception {
// Create a map of consumers and an output target to check output values.
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new
ArrayDeque<>();
- consumers.put(
+ consumers.register(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()),
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<String, Integer>>>)
mainOutputValues::add);
@@ -222,8 +220,7 @@ public class CombineRunnersTest {
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
- FnDataReceiver<WindowedValue<?>> input =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ FnDataReceiver<WindowedValue<?>> input =
consumers.getMultiplexingConsumer(inputPCollectionId);
input.accept(valueInGlobalWindow(KV.of("A", Arrays.asList(1, 2, 6))));
input.accept(valueInGlobalWindow(KV.of("B", Arrays.asList(2, 3))));
input.accept(valueInGlobalWindow(KV.of("C", Arrays.asList(5, 2))));
@@ -243,9 +240,9 @@ public class CombineRunnersTest {
@Test
public void testExtractOutputs() throws Exception {
// Create a map of consumers and an output target to check output values.
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new
ArrayDeque<>();
- consumers.put(
+ consumers.register(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()),
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<String, Integer>>>)
mainOutputValues::add);
@@ -277,8 +274,7 @@ public class CombineRunnersTest {
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
- FnDataReceiver<WindowedValue<?>> input =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ FnDataReceiver<WindowedValue<?>> input =
consumers.getMultiplexingConsumer(inputPCollectionId);
input.accept(valueInGlobalWindow(KV.of("A", 9)));
input.accept(valueInGlobalWindow(KV.of("B", 5)));
input.accept(valueInGlobalWindow(KV.of("C", 7)));
@@ -298,9 +294,9 @@ public class CombineRunnersTest {
@Test
public void testCombineGroupedValues() throws Exception {
// Create a map of consumers and an output target to check output values.
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new
ArrayDeque<>();
- consumers.put(
+ consumers.register(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()),
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<String, Integer>>>)
mainOutputValues::add);
@@ -332,8 +328,7 @@ public class CombineRunnersTest {
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
- FnDataReceiver<WindowedValue<?>> input =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ FnDataReceiver<WindowedValue<?>> input =
consumers.getMultiplexingConsumer(inputPCollectionId);
input.accept(valueInGlobalWindow(KV.of("A", Arrays.asList("1", "2",
"6"))));
input.accept(valueInGlobalWindow(KV.of("B", Arrays.asList("2", "3"))));
input.accept(valueInGlobalWindow(KV.of("C", Arrays.asList("5", "2"))));
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
index 9f2542a..fe661cd 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
@@ -26,16 +26,13 @@ import static org.junit.Assert.assertThat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
-import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Suppliers;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
-import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -68,8 +65,8 @@ public class FlattenRunnerTest {
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
"mainOutputTarget",
(FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
@@ -94,10 +91,10 @@ public class FlattenRunnerTest {
consumers.keySet(),
containsInAnyOrder("inputATarget", "inputBTarget", "inputCTarget",
"mainOutputTarget"));
-
Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A1"));
-
Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A2"));
-
Iterables.getOnlyElement(consumers.get("inputBTarget")).accept(valueInGlobalWindow("B"));
-
Iterables.getOnlyElement(consumers.get("inputCTarget")).accept(valueInGlobalWindow("C"));
+
consumers.getMultiplexingConsumer("inputATarget").accept(valueInGlobalWindow("A1"));
+
consumers.getMultiplexingConsumer("inputATarget").accept(valueInGlobalWindow("A2"));
+
consumers.getMultiplexingConsumer("inputBTarget").accept(valueInGlobalWindow("B"));
+
consumers.getMultiplexingConsumer("inputCTarget").accept(valueInGlobalWindow("C"));
assertThat(
mainOutputValues,
contains(
@@ -131,8 +128,8 @@ public class FlattenRunnerTest {
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
"mainOutputTarget",
(FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
@@ -155,10 +152,9 @@ public class FlattenRunnerTest {
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder("inputATarget",
"mainOutputTarget"));
- assertThat(consumers.get("inputATarget"), hasSize(2));
+ assertThat(consumers.getUnderlyingConsumers("inputATarget"), hasSize(2));
- FnDataReceiver<WindowedValue<?>> input =
- MultiplexingFnDataReceiver.forConsumers(consumers.get("inputATarget"));
+ FnDataReceiver<WindowedValue<?>> input =
consumers.getMultiplexingConsumer("inputATarget");
input.accept(WindowedValue.valueInGlobalWindow("A1"));
input.accept(WindowedValue.valueInGlobalWindow("A2"));
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
index 4017c20..64ca20f 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
@@ -27,13 +27,14 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;
-import com.google.auto.value.AutoValue;
import java.io.Closeable;
import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.HashMap;
import java.util.List;
import java.util.ServiceLoader;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.FakeBeamFnStateClient;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
@@ -42,9 +43,12 @@ import
org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.PipelineTranslation;
import org.apache.beam.runners.core.construction.SdkComponents;
+import org.apache.beam.runners.core.construction.metrics.MetricKey;
import org.apache.beam.runners.core.metrics.MetricUpdates;
import org.apache.beam.runners.core.metrics.MetricUpdates.MetricUpdate;
import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
+import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
@@ -85,10 +89,8 @@ import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Suppliers;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.hamcrest.collection.IsMapContaining;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -196,8 +198,8 @@ public class FnApiDoFnRunnerTest implements Serializable {
bagUserStateKey("combine", "X"), encode("X0")));
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
outputPCollectionId,
(FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
@@ -227,7 +229,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
// Ensure that bag user state that is initially empty or populated works.
// Ensure that the key order does not matter when we traverse over KV
pairs.
FnDataReceiver<WindowedValue<?>> mainInput =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ consumers.getMultiplexingConsumer(inputPCollectionId);
mainInput.accept(valueInGlobalWindow(KV.of("X", "X1")));
mainInput.accept(valueInGlobalWindow(KV.of("Y", "Y1")));
mainInput.accept(valueInGlobalWindow(KV.of("X", "X2")));
@@ -355,11 +357,11 @@ public class FnApiDoFnRunnerTest implements Serializable {
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
outputPCollectionId,
(FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
- consumers.put(
+ consumers.register(
additionalPCollectionId,
(FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
additionalOutputValues::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
@@ -391,7 +393,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
// Ensure that bag user state that is initially empty or populated works.
// Ensure that the bagUserStateKey order does not matter when we traverse
over KV pairs.
FnDataReceiver<WindowedValue<?>> mainInput =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ consumers.getMultiplexingConsumer(inputPCollectionId);
mainInput.accept(valueInGlobalWindow("X"));
mainInput.accept(valueInGlobalWindow("Y"));
assertThat(
@@ -422,8 +424,9 @@ public class FnApiDoFnRunnerTest implements Serializable {
private static class TestSideInputIsAccessibleForDownstreamCallersDoFn
extends DoFn<String, Iterable<String>> {
+ public static final String USER_COUNTER_NAME = "userCountedElems";
private final Counter countedElements =
-
Metrics.counter(TestSideInputIsAccessibleForDownstreamCallersDoFn.class,
"countedElems");
+
Metrics.counter(TestSideInputIsAccessibleForDownstreamCallersDoFn.class,
USER_COUNTER_NAME);
private final PCollectionView<Iterable<String>> iterableSideInput;
@@ -486,8 +489,8 @@ public class FnApiDoFnRunnerTest implements Serializable {
FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient(stateData);
List<WindowedValue<Iterable<String>>> mainOutputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()),
(FnDataReceiver) (FnDataReceiver<WindowedValue<Iterable<String>>>)
mainOutputValues::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
@@ -517,7 +520,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
// Ensure that bag user state that is initially empty or populated works.
// Ensure that the bagUserStateKey order does not matter when we traverse
over KV pairs.
FnDataReceiver<WindowedValue<?>> mainInput =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ consumers.getMultiplexingConsumer(inputPCollectionId);
mainInput.accept(valueInWindow("X", windowA));
mainInput.accept(valueInWindow("Y", windowB));
assertThat(mainOutputValues, hasSize(2));
@@ -532,21 +535,9 @@ public class FnApiDoFnRunnerTest implements Serializable {
assertEquals(stateData, fakeClient.getData());
}
- /**
- * A simple Tuple class for creating a list of ExpectedMetrics using the
stepName, metricName and
- * value of the MetricUpdate classes.
- */
- @AutoValue
- public abstract static class ExpectedMetric implements Serializable {
- static ExpectedMetric create(String stepName, MetricName metricName, long
value) {
- return new AutoValue_FnApiDoFnRunnerTest_ExpectedMetric(stepName,
metricName, value);
- }
-
- public abstract String stepName();
-
- public abstract MetricName metricName();
-
- public abstract long value();
+ /** @return a test MetricUpdate for expected metrics to compare against */
+ public MetricUpdate create(String stepName, MetricName name, long value) {
+ return MetricUpdate.create(MetricKey.create(stepName, name), value);
}
@Test
@@ -575,6 +566,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
SdkComponents sdkComponents = SdkComponents.create(p.getOptions());
RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents,
true);
String inputPCollectionId =
sdkComponents.registerPCollection(valuePCollection);
+ String outputPCollectionId =
sdkComponents.registerPCollection(outputPCollection);
RunnerApi.PTransform pTransform =
pProto
@@ -597,8 +589,8 @@ public class FnApiDoFnRunnerTest implements Serializable {
FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient(stateData);
List<WindowedValue<Iterable<String>>> mainOutputValues = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
Iterables.getOnlyElement(pTransform.getOutputsMap().values()),
(FnDataReceiver) (FnDataReceiver<WindowedValue<Iterable<String>>>)
mainOutputValues::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
@@ -623,31 +615,41 @@ public class FnApiDoFnRunnerTest implements Serializable {
Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
mainOutputValues.clear();
+ assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
+
// Ensure that bag user state that is initially empty or populated works.
// Ensure that the bagUserStateKey order does not matter when we traverse
over KV pairs.
FnDataReceiver<WindowedValue<?>> mainInput =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ consumers.getMultiplexingConsumer(inputPCollectionId);
mainInput.accept(valueInWindow("X", windowA));
mainInput.accept(valueInWindow("Y", windowB));
MetricsContainer mc = MetricsEnvironment.getCurrentContainer();
+
+ List<MetricUpdate> expectedMetrics = new ArrayList<MetricUpdate>();
MetricName metricName =
-
MetricName.named(TestSideInputIsAccessibleForDownstreamCallersDoFn.class,
"countedElems");
- List<ExpectedMetric> expectedMetrics = new ArrayList<ExpectedMetric>();
- expectedMetrics.add(ExpectedMetric.create("testUsingMetrics", metricName,
2));
+ MetricName.named(
+ TestSideInputIsAccessibleForDownstreamCallersDoFn.class,
+
TestSideInputIsAccessibleForDownstreamCallersDoFn.USER_COUNTER_NAME);
+ expectedMetrics.add(create("testUsingMetrics", metricName, 2));
+
+ HashMap<String, String> labels = new HashMap<String, String>();
+
+ labels.put(SimpleMonitoringInfoBuilder.PCOLLECTION_LABEL,
"Window.Into()/Window.Assign.out");
+ metricName =
+
MonitoringInfoMetricName.named(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN,
labels);
+ expectedMetrics.add(create("testUsingMetrics", metricName, 2));
+
+ labels.put(
+ SimpleMonitoringInfoBuilder.PCOLLECTION_LABEL,
+
"pTransformId/ParMultiDo(TestSideInputIsAccessibleForDownstreamCallers).output");
+ metricName =
+
MonitoringInfoMetricName.named(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN,
labels);
+ expectedMetrics.add(create("testUsingMetrics", metricName, 2));
closeable.close();
MetricUpdates updates = metricsContainer.getUpdates();
-
- // Validate MetricUpdates
- int i = 0;
- for (MetricUpdate mu : updates.counterUpdates()) {
- assertEquals(expectedMetrics.get(i).metricName(),
mu.getKey().metricName());
- assertEquals(expectedMetrics.get(i).stepName(), mu.getKey().stepName());
- assertEquals(expectedMetrics.get(i).value(), mu.getUpdate());
- i++;
- }
- assertEquals(1, i); // Validate the length.
+ assertThat(updates.counterUpdates(),
containsInAnyOrder(expectedMetrics.toArray()));
}
private static class TestTimerfulDoFn extends DoFn<KV<String, String>,
String> {
@@ -746,15 +748,15 @@ public class FnApiDoFnRunnerTest implements Serializable {
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
List<WindowedValue<KV<String, Timer>>> eventTimerOutputValues = new
ArrayList<>();
List<WindowedValue<KV<String, Timer>>> processingTimerOutputValues = new
ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put(
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register(
outputPCollectionId,
(FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
- consumers.put(
+ consumers.register(
eventTimerOutputPCollectionId,
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<String, Timer>>>)
eventTimerOutputValues::add);
- consumers.put(
+ consumers.register(
processingTimerOutputPCollectionId,
(FnDataReceiver)
(FnDataReceiver<WindowedValue<KV<String, Timer>>>)
processingTimerOutputValues::add);
@@ -806,11 +808,11 @@ public class FnApiDoFnRunnerTest implements Serializable {
// Ensure that bag user state that is initially empty or populated works.
// Ensure that the key order does not matter when we traverse over KV
pairs.
FnDataReceiver<WindowedValue<?>> mainInput =
- Iterables.getOnlyElement(consumers.get(inputPCollectionId));
+ consumers.getMultiplexingConsumer(inputPCollectionId);
FnDataReceiver<WindowedValue<?>> eventTimerInput =
- Iterables.getOnlyElement(consumers.get(eventTimerInputPCollectionId));
+ consumers.getMultiplexingConsumer(eventTimerInputPCollectionId);
FnDataReceiver<WindowedValue<?>> processingTimerInput =
-
Iterables.getOnlyElement(consumers.get(processingTimerInputPCollectionId));
+ consumers.getMultiplexingConsumer(processingTimerInputPCollectionId);
mainInput.accept(timestampedValueInGlobalWindow(KV.of("X", "X1"), new
Instant(1000L)));
mainInput.accept(timestampedValueInGlobalWindow(KV.of("Y", "Y1"), new
Instant(1100L)));
mainInput.accept(timestampedValueInGlobalWindow(KV.of("X", "X2"), new
Instant(1200L)));
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
index 1aa0eb5..d93cb6a 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
@@ -28,10 +28,10 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.apache.beam.fn.harness.MapFnRunners.ValueMapFnFactory;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.fn.function.ThrowingFunction;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -39,10 +39,7 @@ import
org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Suppliers;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ArrayListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSet;
-import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Test;
@@ -62,8 +59,8 @@ public class MapFnRunnersTest {
@Test
public void testValueOnlyMapping() throws Exception {
List<WindowedValue<?>> outputConsumer = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put("outputPC", outputConsumer::add);
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register("outputPC", outputConsumer::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry = new
PTransformFunctionRegistry();
@@ -90,7 +87,7 @@ public class MapFnRunnersTest {
assertThat(consumers.keySet(), containsInAnyOrder("inputPC", "outputPC"));
-
Iterables.getOnlyElement(consumers.get("inputPC")).accept(valueInGlobalWindow("abc"));
+
consumers.getMultiplexingConsumer("inputPC").accept(valueInGlobalWindow("abc"));
assertThat(outputConsumer, contains(valueInGlobalWindow("ABC")));
}
@@ -98,8 +95,8 @@ public class MapFnRunnersTest {
@Test
public void testFullWindowedValueMapping() throws Exception {
List<WindowedValue<?>> outputConsumer = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put("outputPC", outputConsumer::add);
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register("outputPC", outputConsumer::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry = new
PTransformFunctionRegistry();
@@ -125,7 +122,7 @@ public class MapFnRunnersTest {
assertThat(consumers.keySet(), containsInAnyOrder("inputPC", "outputPC"));
-
Iterables.getOnlyElement(consumers.get("inputPC")).accept(valueInGlobalWindow("abc"));
+
consumers.getMultiplexingConsumer("inputPC").accept(valueInGlobalWindow("abc"));
assertThat(outputConsumer, contains(valueInGlobalWindow("ABC")));
}
@@ -133,8 +130,8 @@ public class MapFnRunnersTest {
@Test
public void testFullWindowedValueMappingWithCompressedWindow() throws
Exception {
List<WindowedValue<?>> outputConsumer = new ArrayList<>();
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>> consumers =
ArrayListMultimap.create();
- consumers.put("outputPC", outputConsumer::add);
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ consumers.register("outputPC", outputConsumer::add);
PTransformFunctionRegistry startFunctionRegistry = new
PTransformFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry = new
PTransformFunctionRegistry();
@@ -163,7 +160,8 @@ public class MapFnRunnersTest {
IntervalWindow firstWindow = new IntervalWindow(new Instant(0L),
Duration.standardMinutes(10L));
IntervalWindow secondWindow =
new IntervalWindow(new Instant(-10L), Duration.standardSeconds(22L));
- Iterables.getOnlyElement(consumers.get("inputPC"))
+ consumers
+ .getMultiplexingConsumer("inputPC")
.accept(
WindowedValue.of(
"abc",
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
index 130ebac..1d82009 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
@@ -34,6 +34,7 @@ import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.PTransformRunnerFactory;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache;
@@ -46,14 +47,12 @@ import org.apache.beam.model.pipeline.v1.RunnerApi.Coder;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowingStrategy;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.fn.function.ThrowingConsumer;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.Message;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
-import
org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ListMultimap;
import
org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.Uninterruptibles;
import org.junit.Before;
import org.junit.Rule;
@@ -116,7 +115,7 @@ public class ProcessBundleHandlerTest {
pCollections,
coders,
windowingStrategies,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
startFunctionRegistry,
finishFunctionRegistry,
splitListener) -> {
@@ -187,7 +186,7 @@ public class ProcessBundleHandlerTest {
pCollections,
coders,
windowingStrategies,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
startFunctionRegistry,
finishFunctionRegistry,
splitListener) -> {
@@ -233,7 +232,7 @@ public class ProcessBundleHandlerTest {
pCollections,
coders,
windowingStrategies,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
startFunctionRegistry,
finishFunctionRegistry,
splitListener) -> {
@@ -281,7 +280,7 @@ public class ProcessBundleHandlerTest {
pCollections,
coders,
windowingStrategies,
- pCollectionIdsToConsumers,
+ pCollectionConsumerRegistry,
startFunctionRegistry,
finishFunctionRegistry,
splitListener) -> {
@@ -367,8 +366,7 @@ public class ProcessBundleHandlerTest {
Map<String, PCollection> pCollections,
Map<String, Coder> coders,
Map<String, WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
- pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
@@ -430,8 +428,7 @@ public class ProcessBundleHandlerTest {
Map<String, PCollection> pCollections,
Map<String, Coder> coders,
Map<String, WindowingStrategy> windowingStrategies,
- ListMultimap<String, FnDataReceiver<WindowedValue<?>>>
- pCollectionIdsToConsumers,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
PTransformFunctionRegistry startFunctionRegistry,
PTransformFunctionRegistry finishFunctionRegistry,
BundleSplitListener splitListener)
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiverTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiverTest.java
new file mode 100644
index 0000000..a18c2e9
--- /dev/null
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiverTest.java
@@ -0,0 +1,70 @@
+/*
+ * 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.fn.harness.data;
+
+import static junit.framework.TestCase.assertEquals;
+import static org.mockito.Mockito.mock;
+
+import java.io.Closeable;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.MonitoringInfo;
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.junit.Test;
+
+/** Tests for {@link ElementCountFnDataReceiver}. */
+public class ElementCountFnDataReceiverTest {
+
+ /**
+ * Test that the elements are counted, and a MonitoringInfo can be extracted
from a
+ * metricsContainer, if it is in scope.
+ *
+ * @throws Exception
+ */
+ @Test
+ public void testCountsElements() throws Exception {
+ final String pCollectionA = "pCollectionA";
+ final String stepA = "stepA";
+
+ MetricsContainerImpl metricsContainer = new MetricsContainerImpl(stepA);
+ try (Closeable closeable =
MetricsEnvironment.scopedMetricsContainer(metricsContainer)) {
+
+ FnDataReceiver<WindowedValue<String>> consumer =
mock(FnDataReceiver.class);
+ ElementCountFnDataReceiver<String> wrapperConsumer =
+ new ElementCountFnDataReceiver(consumer, pCollectionA);
+ WindowedValue<String> element =
WindowedValue.valueInGlobalWindow("elem");
+ int numElements = 20;
+ for (int i = 0; i < numElements; i++) {
+ wrapperConsumer.accept(element);
+ }
+
+ SimpleMonitoringInfoBuilder builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(SimpleMonitoringInfoBuilder.ELEMENT_COUNT_URN);
+ builder.setPCollectionLabel(pCollectionA);
+ builder.setInt64Value(numElements);
+ MonitoringInfo expected = builder.build();
+
+ // Clear the timestamp before compairison.
+ MonitoringInfo first =
metricsContainer.getMonitoringInfos().iterator().next();
+ MonitoringInfo result =
SimpleMonitoringInfoBuilder.clearTimestamp(first);
+ assertEquals(expected, result);
+ }
+ }
+}
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java
new file mode 100644
index 0000000..beba73e
--- /dev/null
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java
@@ -0,0 +1,88 @@
+/*
+ * 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.fn.harness.data;
+
+import static org.hamcrest.Matchers.contains;
+import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link PCollectionConsumerRegistryTest}. */
+@RunWith(JUnit4.class)
+public class PCollectionConsumerRegistryTest {
+
+ @Rule public ExpectedException expectedException = ExpectedException.none();
+
+ /**
+ * Test that the counter increments only once when multiple consumers of
same pCollection read the
+ * same element.
+ */
+ @Test
+ public void multipleConsumersSamePCollection() throws Exception {
+ final String pCollectionA = "pCollectionA";
+
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ FnDataReceiver<WindowedValue<String>> consumerA1 =
mock(FnDataReceiver.class);
+ FnDataReceiver<WindowedValue<String>> consumerA2 =
mock(FnDataReceiver.class);
+
+ consumers.register(pCollectionA, consumerA1);
+ consumers.register(pCollectionA, consumerA2);
+
+ FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+ (FnDataReceiver<WindowedValue<String>>)
+ (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+
+ WindowedValue<String> element = WindowedValue.valueInGlobalWindow("elem");
+ int numElements = 20;
+ for (int i = 0; i < numElements; i++) {
+ wrapperConsumer.accept(element);
+ }
+
+ // Check that the underlying consumers are each invoked per element.
+ verify(consumerA1, times(numElements)).accept(element);
+ verify(consumerA2, times(numElements)).accept(element);
+
+ assertThat(consumers.keySet(), contains(pCollectionA));
+ assertThat(consumers.getUnderlyingConsumers(pCollectionA),
contains(consumerA1, consumerA2));
+ }
+
+ @Test
+ public void throwsOnRegisteringAfterMultiplexingConsumerWasInitialized()
throws Exception {
+ final String pCollectionA = "pCollectionA";
+
+ PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry();
+ FnDataReceiver<WindowedValue<String>> consumerA1 =
mock(FnDataReceiver.class);
+ FnDataReceiver<WindowedValue<String>> consumerA2 =
mock(FnDataReceiver.class);
+
+ consumers.register(pCollectionA, consumerA1);
+ consumers.getMultiplexingConsumer(pCollectionA);
+
+ expectedException.expect(RuntimeException.class);
+ expectedException.expectMessage("cannot be register()-d after");
+ consumers.register(pCollectionA, consumerA2);
+ }
+}
