Add CounterSet#merge(CounterSet)

This method merges the contents of the two counter sets into the called
CounterSet, using the underlying Counter#merge method. All counters that
are present only in one of the CounterSets are added to the called set.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=113225321

Author: tgroh

sdk/src/main/java/com/google/cloud/dataflow/sdk/util/common/CounterSet.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.dataflow.sdk.util.common;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import java.util.AbstractSet;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -129,6 +131,29 @@ public synchronized boolean add(Counter<?> e) {
     return true;
   }
 
+  public synchronized void merge(CounterSet that) {
+    for (Counter<?> theirCounter : that) {
+      Counter<?> myCounter = counters.get(theirCounter.getName());
+      if (myCounter != null) {
+        mergeCounters(myCounter, theirCounter);
+      } else {
+        addCounter(theirCounter);
+      }
+    }
+  }
+
+  private <T> void mergeCounters(Counter<T> mine, Counter<?> theirCounter) {
+    checkArgument(
+        mine.isCompatibleWith(theirCounter),
+        "Can't merge CounterSets containing incompatible counters with the same name: "
+            + "%s (existing) and %s (merged)",
+        mine,
+        theirCounter);
+    @SuppressWarnings("unchecked")
+    Counter<T> theirs = (Counter<T>) theirCounter;
+    mine.merge(theirs);
+  }
+
   /**
    * A nested class that supports adding additional counters into the
    * enclosing CounterSet. This is useful as a mutator, hiding other

sdk/src/test/java/com/google/cloud/dataflow/sdk/util/common/CounterSetTest.java

@@ -19,6 +19,7 @@
 import static com.google.cloud.dataflow.sdk.util.common.Counter.AggregationKind.MAX;
 import static com.google.cloud.dataflow.sdk.util.common.Counter.AggregationKind.SUM;
 import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.equalTo;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThat;
@@ -109,6 +110,70 @@ public void testAddOrReuseWithIncompatibleTypesThrowsException() {
     set.addOrReuseCounter(c1Incompatible);
   }
 
+  @Test
+  public void testMergeWithDifferentNamesAddsAll() {
+    Counter<?> c1 = Counter.longs("c1", SUM);
+    Counter<?> c2 = Counter.ints("c2", MAX);
+
+    set.add(c1);
+    set.add(c2);
+
+    CounterSet newSet = new CounterSet();
+    newSet.merge(set);
+
+    assertThat(newSet, containsInAnyOrder(c1, c2));
+  }
+
+  @SuppressWarnings("unchecked")
+  @Test
+  public void testMergeWithSameNamesMerges() {
+    Counter<Long> c1 = Counter.longs("c1", SUM);
+    Counter<Integer> c2 = Counter.ints("c2", MAX);
+
+    set.add(c1);
+    set.add(c2);
+
+    c1.addValue(3L);
+    c2.addValue(22);
+
+    CounterSet newSet = new CounterSet();
+    Counter<Long> c1Prime = Counter.longs("c1", SUM);
+    Counter<Integer> c2Prime = Counter.ints("c2", MAX);
+
+    c1Prime.addValue(7L);
+    c2Prime.addValue(14);
+
+    newSet.add(c1Prime);
+    newSet.add(c2Prime);
+
+    newSet.merge(set);
+
+    assertThat((Counter<Long>) newSet.getExistingCounter("c1"), equalTo(c1Prime));
+    assertThat((Long) newSet.getExistingCounter("c1").getAggregate(), equalTo(10L));
+
+    assertThat((Counter<Integer>) newSet.getExistingCounter("c2"), equalTo(c2Prime));
+    assertThat((Integer) newSet.getExistingCounter("c2").getAggregate(), equalTo(22));
+  }
+
+  @SuppressWarnings("unchecked")
+  @Test
+  public void testMergeWithIncompatibleTypesThrowsException() {
+    Counter<Long> c1 = Counter.longs("c1", SUM);
+
+    set.add(c1);
+
+    CounterSet newSet = new CounterSet();
+    Counter<Long> c1Prime = Counter.longs("c1", MAX);
+
+    newSet.add(c1Prime);
+
+    thrown.expect(IllegalArgumentException.class);
+    thrown.expectMessage("c1");
+    thrown.expectMessage("incompatible counters with the same name");
+
+    newSet.merge(set);
+  }
+
   @Test
   public void testAddCounterMutatorAddCounterAddsCounter() {
     Counter<?> c1 = Counter.longs("c1", SUM);