Finish Java data sampling impl with tests, adding comments

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java

@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.EnumMap;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.function.Function;
@@ -256,9 +257,12 @@ public static void main(
 
       // Add any graph modifications.
       List<ProcessBundleDescriptorModifier> modifierList = new ArrayList<>();
-      List<String> experimentList = options.as(ExperimentalOptions.class).getExperiments();
+      Optional<List<String>> experimentList =
+          Optional.ofNullable(options.as(ExperimentalOptions.class).getExperiments());
 
-      if (experimentList != null && experimentList.contains(ENABLE_DATA_SAMPLING_EXPERIMENT)) {
+      // If data sampling is enabled, then modify the graph to add any DataSampling Operations.
+      if (experimentList.isPresent()
+          && experimentList.get().contains(ENABLE_DATA_SAMPLING_EXPERIMENT)) {
         modifierList.add(new DataSamplingDescriptorModifier());
       }
 
@@ -357,6 +361,7 @@ private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
       handlers.put(
           InstructionRequest.RequestCase.HARNESS_MONITORING_INFOS,
           processWideHandler::harnessMonitoringInfos);
+      handlers.put(InstructionRequest.RequestCase.SAMPLE, dataSampler::handleDataSampleRequest);
 
       JvmInitializers.runBeforeProcessing(options);
 

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

@@ -17,33 +17,65 @@
  */
 package org.apache.beam.fn.harness.debug;
 
+import avro.shaded.com.google.common.collect.ImmutableSet;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+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.model.fnexecution.v1.BeamFnApi.SampledElement;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.vendor.grpc.v1p48p1.com.google.protobuf.ByteString;
 
+/**
+ * The DataSampler is a global (per SDK Harness) object that facilitates taking and returning
+ * samples to the Runner Harness. The class is thread-safe with respect to executing
+ * ProcessBundleDescriptors. Meaning, different threads executing different PBDs can sample
+ * simultaneously, even if computing the same logical PCollection.
+ */
 public class DataSampler {
 
+  /** Creates a DataSampler to sample every 10 elements while keeping a maximum of 10 in memory. */
   public DataSampler() {}
 
+  /**
+   * @param maxSamples Sets the maximum number of samples held in memory at once.
+   * @param sampleEveryN Sets how often to sample.
+   */
   public DataSampler(int maxSamples, int sampleEveryN) {
     this.maxSamples = maxSamples;
     this.sampleEveryN = sampleEveryN;
   }
 
-  public static Set<String> EMPTY = new HashSet<>();
-
   // Maximum number of elements in buffer.
   private int maxSamples = 10;
 
   // Sampling rate.
   private int sampleEveryN = 1000;
 
-  private final Map<String, Map<String, OutputSampler<?>>> outputSamplers = new HashMap<>();
-
+  // The fully-qualified type is: Map[ProcessBundleDescriptorId, [PCollectionId, OutputSampler]].
+  // The DataSampler object lives on the same level of the FnHarness. This means that many threads
+  // can and will
+  // access this simultaneously. However, ProcessBundleDescriptors are unique per thread, so only
+  // synchronization
+  // is needed on the outermost map.
+  private final Map<String, Map<String, OutputSampler<?>>> outputSamplers =
+      new ConcurrentHashMap<>();
+
+  /**
+   * Creates and returns a class to sample the given PCollection in the given
+   * ProcessBundleDescriptor. Uses the given coder encode samples as bytes when responding to a
+   * SampleDataRequest.
+   *
+   * @param processBundleDescriptorId The PBD to sample from.
+   * @param pcollectionId The PCollection to take intermittent samples from.
+   * @param coder The coder associated with the PCollection. Coder may be from a nested context.
+   * @return the OutputSampler corresponding to the unique PBD and PCollection.
+   * @param <T> The type of element contained in the PCollection.
+   */
   public <T> OutputSampler<T> sampleOutput(
       String processBundleDescriptorId, String pcollectionId, Coder<T> coder) {
     outputSamplers.putIfAbsent(processBundleDescriptorId, new HashMap<>());
@@ -54,12 +86,50 @@ public <T> OutputSampler<T> sampleOutput(
     return (OutputSampler<T>) samplers.get(pcollectionId);
   }
 
-  public Map<String, List<byte[]>> samples() {
-    return samplesFor(EMPTY, EMPTY);
+  /**
+   * Returns all collected samples. Thread-safe.
+   *
+   * @param request The instruction request from the FnApi. Filters based on the given
+   *     SampleDataRequest.
+   * @return Returns all collected samples.
+   */
+  public BeamFnApi.InstructionResponse.Builder handleDataSampleRequest(
+      BeamFnApi.InstructionRequest request) {
+    BeamFnApi.SampleDataRequest sampleDataRequest = request.getSample();
+
+    Map<String, List<byte[]>> responseSamples =
+        samplesFor(
+            ImmutableSet.copyOf(sampleDataRequest.getProcessBundleDescriptorIdsList()),
+            ImmutableSet.copyOf(sampleDataRequest.getPcollectionIdsList()));
+
+    BeamFnApi.SampleDataResponse.Builder response = BeamFnApi.SampleDataResponse.newBuilder();
+    for (String pcollectionId : responseSamples.keySet()) {
+      ElementList.Builder elementList = ElementList.newBuilder();
+      for (byte[] sample : responseSamples.get(pcollectionId)) {
+        elementList.addElements(
+            SampledElement.newBuilder().setElement(ByteString.copyFrom(sample)).build());
+      }
+      response.putElementSamples(pcollectionId, elementList.build());
+    }
+
+    return BeamFnApi.InstructionResponse.newBuilder().setSample(response);
   }
 
+  /**
+   * Returns a map from PCollection to its samples. Samples are filtered on
+   * ProcessBundleDescriptorIds and PCollections. Thread-safe.
+   *
+   * @param descriptors PCollections under each PBD id will be unioned. If empty, allows all
+   *     descriptors.
+   * @param pcollections Filters all PCollections on this set. If empty, allows all PCollections.
+   * @return a map from PCollection to its samples.
+   */
   public Map<String, List<byte[]>> samplesFor(Set<String> descriptors, Set<String> pcollections) {
     Map<String, List<byte[]>> samples = new HashMap<>();
+
+    // Safe to iterate as the ConcurrentHashMap will return each element at most once and will not
+    // throw
+    // ConcurrentModificationException.
     outputSamplers.forEach(
         (descriptorId, samplers) -> {
           if (!descriptors.isEmpty() && !descriptors.contains(descriptorId)) {
@@ -80,11 +150,24 @@ public Map<String, List<byte[]>> samplesFor(Set<String> descriptors, Set<String>
     return samples;
   }
 
+  /** @return samples from all PBDs and all PCollections. */
+  public Map<String, List<byte[]>> allSamples() {
+    return samplesFor(ImmutableSet.of(), ImmutableSet.of());
+  }
+
+  /**
+   * @param descriptors PBDs to filter on.
+   * @return samples only from the given descriptors.
+   */
   public Map<String, List<byte[]>> samplesForDescriptors(Set<String> descriptors) {
-    return samplesFor(descriptors, EMPTY);
+    return samplesFor(descriptors, ImmutableSet.of());
   }
 
+  /**
+   * @param pcollections PCollection ids to filter on.
+   * @return samples only from the given PCollections.
+   */
   public Map<String, List<byte[]>> samplesForPCollections(Set<String> pcollections) {
-    return samplesFor(EMPTY, pcollections);
+    return samplesFor(ImmutableSet.of(), pcollections);
   }
 }

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

@@ -27,8 +27,10 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.*;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.vendor.grpc.v1p48p1.com.google.protobuf.ByteString;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -49,17 +51,27 @@ public byte[] encodeString(String s) throws IOException {
     return stream.toByteArray();
   }
 
+  /**
+   * Smoke test that a samples show in the output map.
+   *
+   * @throws Exception
+   */
   @Test
   public void testSingleOutput() throws Exception {
     DataSampler sampler = new DataSampler();
 
     VarIntCoder coder = VarIntCoder.of();
     sampler.sampleOutput("descriptor-id", "pcollection-id", coder).sample(1);
 
-    Map<String, List<byte[]>> samples = sampler.samples();
+    Map<String, List<byte[]>> samples = sampler.allSamples();
     assertThat(samples.get("pcollection-id"), contains(encodeInt(1)));
   }
 
+  /**
+   * Test that sampling multiple PCollections under the same descriptor is OK.
+   *
+   * @throws Exception
+   */
   @Test
   public void testMultipleOutputs() throws Exception {
     DataSampler sampler = new DataSampler();
@@ -68,11 +80,16 @@ public void testMultipleOutputs() throws Exception {
     sampler.sampleOutput("descriptor-id", "pcollection-id-1", coder).sample(1);
     sampler.sampleOutput("descriptor-id", "pcollection-id-2", coder).sample(2);
 
-    Map<String, List<byte[]>> samples = sampler.samples();
+    Map<String, List<byte[]>> samples = sampler.allSamples();
     assertThat(samples.get("pcollection-id-1"), contains(encodeInt(1)));
     assertThat(samples.get("pcollection-id-2"), contains(encodeInt(2)));
   }
 
+  /**
+   * Test that the response contains samples from the same PCollection across descriptors.
+   *
+   * @throws Exception
+   */
   @Test
   public void testMultipleDescriptors() throws Exception {
     DataSampler sampler = new DataSampler();
@@ -81,10 +98,15 @@ public void testMultipleDescriptors() throws Exception {
     sampler.sampleOutput("descriptor-id-1", "pcollection-id", coder).sample(1);
     sampler.sampleOutput("descriptor-id-2", "pcollection-id", coder).sample(2);
 
-    Map<String, List<byte[]>> samples = sampler.samples();
+    Map<String, List<byte[]>> samples = sampler.allSamples();
     assertThat(samples.get("pcollection-id"), contains(encodeInt(1), encodeInt(2)));
   }
 
+  /**
+   * Test that samples can be filtered based on ProcessBundleDescriptor id.
+   *
+   * @throws Exception
+   */
   @Test
   public void testFiltersSingleDescriptorId() throws Exception {
     DataSampler sampler = new DataSampler(10, 10);
@@ -101,6 +123,11 @@ public void testFiltersSingleDescriptorId() throws Exception {
     assertThat(samples.get("2"), contains(encodeString("a2")));
   }
 
+  /**
+   * Test that samples are unioned based on ProcessBundleDescriptor id.
+   *
+   * @throws Exception
+   */
   @Test
   public void testFiltersMultipleDescriptorId() throws Exception {
     DataSampler sampler = new DataSampler(10, 10);
@@ -116,6 +143,11 @@ public void testFiltersMultipleDescriptorId() throws Exception {
     assertThat(samples.get("2"), contains(encodeString("a2"), encodeString("b2")));
   }
 
+  /**
+   * Test that samples can be filtered based on PCollection id.
+   *
+   * @throws Exception
+   */
   @Test
   public void testFiltersSinglePCollectionId() throws Exception {
     DataSampler sampler = new DataSampler(10, 10);
@@ -131,14 +163,6 @@ public void testFiltersSinglePCollectionId() throws Exception {
     assertThat(samples.get("1"), containsInAnyOrder(encodeString("a1"), encodeString("b1")));
   }
 
-  Map<String, List<byte[]>> singletonSample(String pcollectionId, byte[] element) {
-    Map<String, List<byte[]>> ret = new HashMap<>();
-    List<byte[]> list = new ArrayList<>();
-    list.add(element);
-    ret.put(pcollectionId, list);
-    return ret;
-  }
-
   void generateStringSamples(DataSampler sampler) {
     StringUtf8Coder coder = StringUtf8Coder.of();
     sampler.sampleOutput("a", "1", coder).sample("a1");
@@ -147,11 +171,17 @@ void generateStringSamples(DataSampler sampler) {
     sampler.sampleOutput("b", "2", coder).sample("b2");
   }
 
+  /**
+   * Test that samples can be filtered both on PCollection and ProcessBundleDescriptor id.
+   *
+   * @throws Exception
+   */
   @Test
   public void testFiltersDescriptorAndPCollectionIds() throws Exception {
     List<String> descriptorIds = ImmutableList.of("a", "b");
     List<String> pcollectionIds = ImmutableList.of("1", "2");
 
+    // Try all combinations for descriptor and PCollection ids.
     for (String descriptorId : descriptorIds) {
       for (String pcollectionId : pcollectionIds) {
         DataSampler sampler = new DataSampler(10, 10);
@@ -166,4 +196,46 @@ public void testFiltersDescriptorAndPCollectionIds() throws Exception {
       }
     }
   }
+
+  /**
+   * Test that the DataSampler can respond with the correct samples with filters.
+   *
+   * @throws Exception
+   */
+  @Test
+  public void testMakesCorrectResponse() throws Exception {
+    DataSampler dataSampler = new DataSampler();
+    generateStringSamples(dataSampler);
+
+    // SampleDataRequest that filters on PCollection=1 and PBD ids = "a" or "b".
+    BeamFnApi.InstructionRequest request =
+        BeamFnApi.InstructionRequest.newBuilder()
+            .setSample(
+                BeamFnApi.SampleDataRequest.newBuilder()
+                    .addPcollectionIds("1")
+                    .addProcessBundleDescriptorIds("a")
+                    .addProcessBundleDescriptorIds("b")
+                    .build())
+            .build();
+    BeamFnApi.InstructionResponse actual = dataSampler.handleDataSampleRequest(request).build();
+    BeamFnApi.InstructionResponse expected =
+        BeamFnApi.InstructionResponse.newBuilder()
+            .setSample(
+                BeamFnApi.SampleDataResponse.newBuilder()
+                    .putElementSamples(
+                        "1",
+                        BeamFnApi.SampleDataResponse.ElementList.newBuilder()
+                            .addElements(
+                                BeamFnApi.SampledElement.newBuilder()
+                                    .setElement(ByteString.copyFrom(encodeString("a1")))
+                                    .build())
+                            .addElements(
+                                BeamFnApi.SampledElement.newBuilder()
+                                    .setElement(ByteString.copyFrom(encodeString("b1")))
+                                    .build())
+                            .build())
+                    .build())
+            .build();
+    assertThat(actual, equalTo(expected));
+  }
 }

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

@@ -89,7 +89,7 @@ public void testCreatingAndProcessingWithSampling() throws Exception {
             .withPipeline(Pipeline.create());
     Coder<String> rehydratedCoder = (Coder<String>) rehydratedComponents.getCoder("coder-id");
 
-    Map<String, List<byte[]>> samples = dataSampler.samples();
+    Map<String, List<byte[]>> samples = dataSampler.allSamples();
     assertThat(samples.keySet(), contains("inputTarget"));
 
     // Ensure that the value was sampled.

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

@@ -56,7 +56,7 @@ public void testSamplesFirstN() throws Exception {
       expected.add(encodeInt(i));
     }
 
-    Map<String, List<byte[]>> samples = sampler.samples();
+    Map<String, List<byte[]>> samples = sampler.allSamples();
     assertThat(samples.get("pcollection-id"), containsInAnyOrder(expected.toArray()));
   }
 
@@ -79,7 +79,7 @@ public void testActsLikeCircularBuffer() throws Exception {
     expected.add(encodeInt(79));
     expected.add(encodeInt(99));
 
-    Map<String, List<byte[]>> samples = sampler.samples();
+    Map<String, List<byte[]>> samples = sampler.allSamples();
     assertThat(samples.get("pcollection-id"), containsInAnyOrder(expected.toArray()));
   }
 }