add concurrency tests

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/debug/DataSampler.java

@@ -17,12 +17,15 @@
  */
 package org.apache.beam.fn.harness.debug;
 
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.SampleDataResponse.ElementList;
 import org.apache.beam.sdk.coders.Coder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * The DataSampler is a global (per SDK Harness) object that facilitates taking and returning
@@ -31,6 +34,7 @@
  * simultaneously, even if computing the same logical PCollection.
  */
 public class DataSampler {
+  private static final Logger LOG = LoggerFactory.getLogger(DataSampler.class);
 
   /**
    * Creates a DataSampler to sample every 1000 elements while keeping a maximum of 10 in memory.
@@ -99,9 +103,13 @@ public synchronized BeamFnApi.InstructionResponse.Builder handleDataSampleReques
             return;
           }
 
-          response.putElementSamples(
-              pcollectionId,
-              ElementList.newBuilder().addAllElements(outputSampler.samples()).build());
+          try {
+            response.putElementSamples(
+                pcollectionId,
+                ElementList.newBuilder().addAllElements(outputSampler.samples()).build());
+          } catch (IOException e) {
+            LOG.warn("Could not encode elements from \"" + pcollectionId + "\" to bytes: " + e);
+          }
         });
 
     return BeamFnApi.InstructionResponse.newBuilder().setSampleData(response);

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/debug/OutputSampler.java

@@ -17,14 +17,13 @@
  */
 package org.apache.beam.fn.harness.debug;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicLong;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.util.ByteStringOutputStream;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * This class holds samples for a single PCollection until queried by the parent DataSampler. This
@@ -34,7 +33,6 @@
  * @param <T> the element type of the PCollection.
  */
 public class OutputSampler<T> {
-  private static final Logger LOG = LoggerFactory.getLogger(OutputSampler.class);
 
   // Temporarily holds elements until the SDK receives a sample data request.
   private final List<T> buffer;
@@ -101,7 +99,7 @@ public void sample(T element) {
    *
    * @return samples taken since last call.
    */
-  public List<BeamFnApi.SampledElement> samples() {
+  public List<BeamFnApi.SampledElement> samples() throws IOException {
     List<BeamFnApi.SampledElement> ret = new ArrayList<>();
 
     // Serializing can take a lot of CPU time for larger or complex elements. Copy the array here
@@ -115,18 +113,12 @@ public List<BeamFnApi.SampledElement> samples() {
 
     ByteStringOutputStream stream = new ByteStringOutputStream();
     for (T el : copiedBuffer) {
-      try {
-        // This is deprecated, but until this is fully removed, this specifically needs the nested
-        // context. This is because the SDK will need to decode the sampled elements with the
-        // ToStringFn.
-        coder.encode(el, stream, Coder.Context.NESTED);
-        ret.add(
-            BeamFnApi.SampledElement.newBuilder()
-                .setElement(stream.toByteStringAndReset())
-                .build());
-      } catch (Exception exception) {
-        LOG.warn("Could not encode element \"" + el + "\" to bytes: " + exception);
-      }
+      // This is deprecated, but until this is fully removed, this specifically needs the nested
+      // context. This is because the SDK will need to decode the sampled elements with the
+      // ToStringFn.
+      coder.encode(el, stream, Coder.Context.NESTED);
+      ret.add(
+          BeamFnApi.SampledElement.newBuilder().setElement(stream.toByteStringAndReset()).build());
     }
 
     return ret;

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java

@@ -21,6 +21,8 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
@@ -39,6 +41,8 @@
 import org.apache.beam.fn.harness.control.BundleProgressReporter;
 import org.apache.beam.fn.harness.control.ExecutionStateSampler;
 import org.apache.beam.fn.harness.control.ExecutionStateSampler.ExecutionStateTracker;
+import org.apache.beam.fn.harness.debug.DataSampler;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor;
 import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
 import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
@@ -56,6 +60,7 @@
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.util.ByteStringOutputStream;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.util.common.ElementByteSizeObservableIterable;
 import org.apache.beam.sdk.util.common.ElementByteSizeObservableIterator;
@@ -507,6 +512,61 @@ public void testLazyByteSizeEstimation() throws Exception {
     assertThat(result, containsInAnyOrder(expected.toArray()));
   }
 
+  /**
+   * Test that element samples are taken when a DataSampler is present.
+   *
+   * @throws Exception
+   */
+  @Test
+  public void dataSampling() throws Exception {
+    final String pTransformIdA = "pTransformIdA";
+
+    ShortIdMap shortIds = new ShortIdMap();
+    BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory();
+    DataSampler dataSampler = new DataSampler();
+    PCollectionConsumerRegistry consumers =
+        new PCollectionConsumerRegistry(
+            sampler.create(), shortIds, reporterAndRegistrar, TEST_DESCRIPTOR, dataSampler);
+    FnDataReceiver<WindowedValue<String>> consumerA1 = mock(FnDataReceiver.class);
+
+    consumers.register(P_COLLECTION_A, pTransformIdA, pTransformIdA + "Name", consumerA1);
+
+    FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+        (FnDataReceiver<WindowedValue<String>>)
+            (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A);
+    String elementValue = "elem";
+    WindowedValue<String> element = valueInGlobalWindow(elementValue);
+    int numElements = 10;
+    for (int i = 0; i < numElements; i++) {
+      wrapperConsumer.accept(element);
+    }
+
+    BeamFnApi.InstructionRequest request =
+        BeamFnApi.InstructionRequest.newBuilder()
+            .setSampleData(BeamFnApi.SampleDataRequest.newBuilder())
+            .build();
+    BeamFnApi.InstructionResponse response = dataSampler.handleDataSampleRequest(request).build();
+
+    Map<String, BeamFnApi.SampleDataResponse.ElementList> elementSamplesMap =
+        response.getSampleData().getElementSamplesMap();
+
+    assertFalse(elementSamplesMap.isEmpty());
+
+    BeamFnApi.SampleDataResponse.ElementList elementList = elementSamplesMap.get(P_COLLECTION_A);
+    assertNotNull(elementList);
+
+    List<BeamFnApi.SampledElement> expectedSamples = new ArrayList<>();
+    StringUtf8Coder coder = StringUtf8Coder.of();
+    for (int i = 0; i < numElements; i++) {
+      ByteStringOutputStream stream = new ByteStringOutputStream();
+      coder.encode(elementValue, stream);
+      expectedSamples.add(
+          BeamFnApi.SampledElement.newBuilder().setElement(stream.toByteStringAndReset()).build());
+    }
+
+    assertTrue(elementList.getElementsList().containsAll(expectedSamples));
+  }
+
   private static class TestElementByteSizeObservableIterable<T>
       extends ElementByteSizeObservableIterable<T, ElementByteSizeObservableIterator<T>> {
     private List<T> elements;

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/debug/DataSamplerTest.java

@@ -221,4 +221,44 @@ public void testFiltersMultiplePCollectionIds() throws Exception {
     assertHasSamples(samples, "a", ImmutableList.of(encodeString("a1"), encodeString("a2")));
     assertHasSamples(samples, "c", ImmutableList.of(encodeString("c1"), encodeString("c2")));
   }
+
+  /**
+   * Test that samples can be taken from the DataSampler while adding new OutputSamplers. This fails
+   * with a ConcurrentModificationException if there is a bug.
+   *
+   * @throws Exception
+   */
+  @Test
+  public void testConcurrentNewSampler() throws Exception {
+    DataSampler sampler = new DataSampler();
+    VarIntCoder coder = VarIntCoder.of();
+
+    // Create a thread that constantly creates new samplers.
+    Thread sampleThread =
+        new Thread(
+            () -> {
+              for (int i = 0; i < 1000000; i++) {
+                sampler.sampleOutput("pcollection-" + i, coder).sample(0);
+
+                // This sleep is here to allow for the test to stop this thread.
+                try {
+                  Thread.sleep(0);
+                } catch (InterruptedException e) {
+                  return;
+                }
+              }
+            });
+
+    sampleThread.start();
+
+    for (int i = 0; i < 20; i++) {
+      sampler.handleDataSampleRequest(
+          BeamFnApi.InstructionRequest.newBuilder()
+              .setSampleData(BeamFnApi.SampleDataRequest.newBuilder())
+              .build());
+    }
+
+    sampleThread.interrupt();
+    sampleThread.join();
+  }
 }

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/debug/OutputSamplerTest.java

@@ -94,4 +94,40 @@ public void testActsLikeCircularBuffer() throws Exception {
     List<BeamFnApi.SampledElement> samples = outputSampler.samples();
     assertThat(samples, containsInAnyOrder(expected.toArray()));
   }
+
+  /**
+   * Test that sampling a PCollection while retrieving samples from multiple threads is ok.
+   *
+   * @throws Exception
+   */
+  @Test
+  public void testConcurrentSamples() throws Exception {
+    VarIntCoder coder = VarIntCoder.of();
+    OutputSampler<Integer> outputSampler = new OutputSampler<>(coder, 100000, 1);
+
+    // Iteration count was empirically chosen to have a high probability of failure without the
+    // test going for too long.
+    Thread sampleThreadA =
+        new Thread(
+            () -> {
+              for (int i = 0; i < 10000000; i++) {
+                outputSampler.sample(i);
+              }
+            });
+
+    Thread sampleThreadB =
+        new Thread(
+            () -> {
+              for (int i = 0; i < 10000000; i++) {
+                outputSampler.sample(i);
+              }
+            });
+
+    sampleThreadA.start();
+    sampleThreadB.start();
+
+    for (int i = 0; i < 10000; i++) {
+      outputSampler.samples();
+    }
+  }
 }