[BEAM-1672] Make MetricsContainers accumulable.

runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataUtils.java

@@ -71,7 +71,7 @@ public boolean apply(@Nullable WindowedValue<V> input) {
  .isBefore(timerInternals.currentInputWatermarkTime());
  if (expired) {
  // The element is too late for this window.
- droppedDueToLateness.inc();
+ droppedDueToLateness.update(1L);
  WindowTracing.debug(
  "GroupAlsoByWindow: Dropping element at {} for key: {}; "
  + "window: {} since it is too far behind inputWatermark: {}",

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/CounterCell.java

@@ -30,7 +30,7 @@
  * indirection.
  */
 @Experimental(Kind.METRICS)
-public class CounterCell implements MetricCell<Long> {
+public class CounterCell implements MetricCell<Counter, Long> {
 
  private final DirtyState dirty = new DirtyState();
  private final AtomicLong value = new AtomicLong();
@@ -41,12 +41,25 @@ public class CounterCell implements MetricCell<Long> {
  */
  CounterCell() {}
 
- /** Increment the counter by the given amount. */
- private void add(long n) {
+ /**
+ * Increment the counter by the given amount.
+ * @param n value to increment by. Can be negative to decrement.
+ */
+ public void update(long n) {
  value.addAndGet(n);
  dirty.afterModification();
  }
 
+ @Override
+ public void update(Long n) {
+ throw new UnsupportedOperationException("CounterCell.update(Long n) should not be used"
+ + " as it performs unnecessary boxing/unboxing. Use CounterCell.update(long n) instead.");
+ }
+
+ @Override public void update(MetricCell<Counter, Long> other) {
+ update((long) other.getCumulative());
+ }
+
  @Override
  public DirtyState getDirty() {
  return dirty;
@@ -56,12 +69,4 @@ public DirtyState getDirty() {
  public Long getCumulative() {
  return value.get();
  }
-
- public void inc() {
- add(1);
- }
-
- public void inc(long n) {
- add(n);
- }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/DirtyState.java

@@ -18,6 +18,7 @@
 
 package org.apache.beam.sdk.metrics;
 
+import java.io.Serializable;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
@@ -41,7 +42,7 @@
  * completed.
  */
 @Experimental(Kind.METRICS)
-class DirtyState {
+class DirtyState implements Serializable {
  private enum State {
  /** Indicates that there have been changes to the MetricCell since last commit. */
  DIRTY,

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/DistributionCell.java

@@ -30,7 +30,7 @@
  * of indirection.
  */
 @Experimental(Kind.METRICS)
-public class DistributionCell implements MetricCell<DistributionData> {
+public class DistributionCell implements MetricCell<Distribution, DistributionData> {
 
  private final DirtyState dirty = new DirtyState();
  private final AtomicReference<DistributionData> value =
@@ -42,15 +42,25 @@ public class DistributionCell implements MetricCell<DistributionData> {
  */
  DistributionCell() {}
 
- /** Increment the counter by the given amount. */
+ /** Increment the distribution by the given amount. */
  public void update(long n) {
+ update(DistributionData.singleton(n));
+ }
+
+ @Override
+ public void update(DistributionData data) {
  DistributionData original;
  do {
  original = value.get();
- } while (!value.compareAndSet(original, original.combine(DistributionData.singleton(n))));
+ } while (!value.compareAndSet(original, original.combine(data)));
  dirty.afterModification();
  }
 
+ @Override
+ public void update(MetricCell<Distribution, DistributionData> other) {
+ update(other.getCumulative());
+ }
+
  @Override
  public DirtyState getDirty() {
  return dirty;

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/GaugeCell.java

@@ -29,17 +29,33 @@
  * of indirection.
  */
 @Experimental(Experimental.Kind.METRICS)
-public class GaugeCell implements MetricCell<GaugeData> {
+public class GaugeCell implements MetricCell<Gauge, GaugeData> {
 
  private final DirtyState dirty = new DirtyState();
  private final AtomicReference<GaugeData> gaugeValue = new AtomicReference<>(GaugeData.empty());
 
+ /** Set the gauge to the given value. */
  public void set(long value) {
+ update(GaugeData.create(value));
+ }
+
+ @Override
+ public void update(GaugeData data) {
+ GaugeData original;
+ do {
+ original = gaugeValue.get();
+ } while (!gaugeValue.compareAndSet(original, original.combine(data)));
+ dirty.afterModification();
+ }
+
+ @Override
+ public void update(MetricCell<Gauge, GaugeData> other) {
  GaugeData original;
  do {
  original = gaugeValue.get();
- } while (!gaugeValue.compareAndSet(original, original.combine(GaugeData.create(value))));
+ } while (!gaugeValue.compareAndSet(original, original.combine(other.getCumulative())));
  dirty.afterModification();
+ update(other.getCumulative());
  }
 
  @Override

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricCell.java

@@ -17,17 +17,29 @@
  */
 package org.apache.beam.sdk.metrics;
 
+import java.io.Serializable;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
 
 /**
  * A {@link MetricCell} is used for accumulating in-memory changes to a metric. It represents a
  * specific metric name in a single context.
  *
+ * @param <UserT> The type of the user interface for reporting changes to this cell.
  * @param <DataT> The type of metric data stored (and extracted) from this cell.
  */
 @Experimental(Kind.METRICS)
-public interface MetricCell<DataT> {
+public interface MetricCell<UserT extends Metric, DataT> extends Serializable {
+
+ /**
+ * Update value of this cell.
+ */
+ void update(DataT data);
+
+ /**
+ * Update value of this cell by merging the value of another cell.
+ */
+ void update(MetricCell<UserT, DataT> other);
 
  /**
  * Return the {@link DirtyState} tracking whether this metric cell contains uncommitted changes.

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java

@@ -118,7 +118,7 @@ private DelegatingCounter(MetricName name) {
  @Override public void inc(long n) {
  MetricsContainer container = MetricsEnvironment.getCurrentContainer();
  if (container != null) {
- container.getCounter(name).inc(n);
+ container.getCounter(name).update(n);
  }
  }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsContainer.java

@@ -20,6 +20,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.collect.ImmutableList;
+import java.io.Serializable;
 import java.util.Map;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
@@ -37,7 +38,7 @@
  * cumulative values/updates.
  */
 @Experimental(Kind.METRICS)
-public class MetricsContainer {
+public class MetricsContainer implements Serializable {
 
  private final String stepName;
 
@@ -96,7 +97,7 @@ public GaugeCell getGauge(MetricName metricName) {
  return gauges.get(metricName);
  }
 
- private <UpdateT, CellT extends MetricCell<UpdateT>>
+ private <UserT extends Metric, UpdateT, CellT extends MetricCell<UserT, UpdateT>>
  ImmutableList<MetricUpdate<UpdateT>> extractUpdates(
  MetricsMap<MetricName, CellT> cells) {
  ImmutableList.Builder<MetricUpdate<UpdateT>> updates = ImmutableList.builder();
@@ -120,8 +121,8 @@ public MetricUpdates getUpdates() {
  extractUpdates(gauges));
  }
 
- private void commitUpdates(MetricsMap<MetricName, ? extends MetricCell<?>> cells) {
- for (MetricCell<?> cell : cells.values()) {
+ private void commitUpdates(MetricsMap<MetricName, ? extends MetricCell<?, ?>> cells) {
+ for (MetricCell<?, ?> cell : cells.values()) {
  cell.getDirty().afterCommit();
  }
  }
@@ -133,9 +134,10 @@ private void commitUpdates(MetricsMap<MetricName, ? extends MetricCell<?>> cells
  public void commitUpdates() {
  commitUpdates(counters);
  commitUpdates(distributions);
+ commitUpdates(gauges);
  }
 
- private <UpdateT, CellT extends MetricCell<UpdateT>>
+ private <UserT extends Metric, UpdateT, CellT extends MetricCell<UserT, UpdateT>>
  ImmutableList<MetricUpdate<UpdateT>> extractCumulatives(
  MetricsMap<MetricName, CellT> cells) {
  ImmutableList.Builder<MetricUpdate<UpdateT>> updates = ImmutableList.builder();
@@ -156,4 +158,21 @@ public MetricUpdates getCumulative() {
  extractCumulatives(distributions),
  extractCumulatives(gauges));
  }
+
+ /**
+ * Update values of this {@link MetricsContainer} by merging the value of another cell.
+ */
+ public void update(MetricsContainer other) {
+ updateCells(counters, other.counters);
+ updateCells(distributions, other.distributions);
+ updateCells(gauges, other.gauges);
+ }
+
+ private <UserT extends Metric, DataT, CellT extends MetricCell<UserT, DataT>> void updateCells(
+ MetricsMap<MetricName, CellT> current,
+ MetricsMap<MetricName, CellT> updates) {
+ for (Map.Entry<MetricName, CellT> counter : updates.entries()) {
+ current.get(counter.getKey()).update(counter.getValue());
+ }
+ }
 }