Multi-threaded test in MetricsContainerImplTest and some clean up

Author: rohitsinha54

runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/BoundedTrieCell.java

@@ -31,7 +31,6 @@
  * In that case retrieving the underlying cell and reporting directly to it avoids a step of
  * indirection.
  */
-// TODO: Write multi-threaded test in MetricContainerImp for this Cell class too.
 public class BoundedTrieCell implements BoundedTrie, MetricCell<BoundedTrieData> {
 
   private final DirtyState dirty = new DirtyState();
@@ -62,6 +61,10 @@ public DirtyState getDirty() {
     return dirty;
   }
 
+  /**
+   * @return Returns a deep copy of the {@link BoundedTrieData} contained in this {@link
+   *     BoundedTrieCell}.
+   */
   @Override
   public synchronized BoundedTrieData getCumulative() {
     return value.getCumulative();

runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/BoundedTrieData.java

@@ -40,7 +40,7 @@
  * <p>Represents data stored in a bounded trie. This data structure is used to efficiently store and
  * aggregate a collection of string sequences, paths/FQN with a limited size.
  *
- * <p>This class is thread-safe but the underlying BoundedTrieNode contained on it isn't. This is
+ * <p>This class is thread-safe but the underlying BoundedTrieNode contained in it isn't. This is
  * intentional for performance concerns. Hence, this class does not expose the contained node and
  * should not be modified to do so in future when used with multiple threads. This class choose to
  * achieve thread-safety through locks rather than just creating and returning immutable instances
@@ -147,8 +147,8 @@ public synchronized BoundedTrieData getCumulative() {
   }
 
   /**
-   * Returns an immutable set of lists, where each list represents a path in the bounded trie. The
-   * last element in each path is a boolean in string representation denoting whether this path was
+   * Returns {@link BoundedTrieResult}, which represents all path in the bounded trie. The last
+   * element in each path is a boolean in string representation denoting whether this path was
    * truncated. i.e. <["a", "b", "false"], ["c", "true"]>
    *
    * @return The set of paths.
@@ -281,7 +281,6 @@ public final String toString() {
   }
 
   // ------------------------------ BoundedTrieNode Implementation ------------------------------
-
   /**
    * BoundedTrieNode implementation. This class is not thread-safe and relies on the {@link
    * BoundedTrieData} which uses this class to ensure thread-safety by acquiring a lock on the root

runners/core-java/src/test/java/org/apache/beam/runners/core/metrics/MetricsContainerImplTest.java

@@ -34,6 +34,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
 import org.apache.beam.sdk.metrics.MetricName;
 import org.apache.beam.sdk.util.HistogramData;
@@ -561,4 +565,32 @@ public void testMatchMetric() {
             Collections.singletonMap("name", "counter"));
     assertFalse(MetricsContainerImpl.matchMetric(elementCountName, allowedMetricUrns));
   }
+
+  @Test
+  public void testBoundedTrieMultithreaded() throws InterruptedException {
+    MetricsContainerImpl container = new MetricsContainerImpl("step1");
+    BoundedTrieCell boundedTrieCell =
+        container.getBoundedTrie(MetricName.named("test", "boundedTrie"));
+    int num_threads = 10;
+    int num_updates_per_thread = 9; // be under the default bound of 100
+
+    CountDownLatch latch = new CountDownLatch(num_threads);
+    ExecutorService executor = Executors.newFixedThreadPool(num_threads);
+    List<Runnable> tasks = new ArrayList<>();
+    for (int i = 0; i < num_threads; i++) {
+      tasks.add(
+          () -> {
+            for (int j = 0; j < num_updates_per_thread; j++) {
+              boundedTrieCell.add("value-" + Thread.currentThread().getId() + "-" + j);
+            }
+            latch.countDown();
+          });
+    }
+
+    tasks.forEach(executor::execute);
+    latch.await(1, TimeUnit.MINUTES);
+    executor.shutdown();
+
+    assertEquals(num_threads * num_updates_per_thread, boundedTrieCell.getCumulative().size());
+  }
 }

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowMetrics.java

@@ -59,7 +59,8 @@
 class DataflowMetrics extends MetricResults {
 
   private static final Logger LOG = LoggerFactory.getLogger(DataflowMetrics.class);
-  // TODO: Remove this
+  // TODO (rosinha): Remove this once bounded_trie is available in metrics proto Dataflow
+  //  java client.
   public static final String BOUNDED_TRIE = "bounded_trie";
   /**
    * Client for the Dataflow service. This can be used to query the service for information about