address pr comments

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

@@ -21,14 +21,12 @@
 import java.util.HashMap;
 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;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
 
 /**
  * The DataSampler is a global (per SDK Harness) object that facilitates taking and returning
@@ -38,8 +36,13 @@
  */
 public class DataSampler {
 
-  /** Creates a DataSampler to sample every 10 elements while keeping a maximum of 10 in memory. */
-  public DataSampler() {}
+  /**
+   * Creates a DataSampler to sample every 1000 elements while keeping a maximum of 10 in memory.
+   */
+  public DataSampler() {
+    this.maxSamples = 10;
+    this.sampleEveryN = 1000;
+  }
 
   /**
    * @param maxSamples Sets the maximum number of samples held in memory at once.
@@ -51,17 +54,15 @@ public DataSampler(int maxSamples, int sampleEveryN) {
   }
 
   // Maximum number of elements in buffer.
-  private final int maxSamples = 10;
+  private final int maxSamples;
 
   // Sampling rate.
-  private final int sampleEveryN = 1000;
+  private final int sampleEveryN;
 
   // 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.
+  // 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<>();
 
@@ -73,8 +74,8 @@ public DataSampler(int maxSamples, int sampleEveryN) {
    * @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.
+   * @return the OutputSampler corresponding to the unique PBD and PCollection.
    */
   public <T> OutputSampler<T> sampleOutput(
       String processBundleDescriptorId, String pcollectionId, Coder<T> coder) {
@@ -99,8 +100,8 @@ public BeamFnApi.InstructionResponse.Builder handleDataSampleRequest(
 
     Map<String, List<byte[]>> responseSamples =
         samplesFor(
-            ImmutableSet.copyOf(sampleDataRequest.getProcessBundleDescriptorIdsList()),
-            ImmutableSet.copyOf(sampleDataRequest.getPcollectionIdsList()));
+            sampleDataRequest.getProcessBundleDescriptorIdsList(),
+            sampleDataRequest.getPcollectionIdsList());
 
     BeamFnApi.SampleDataResponse.Builder response = BeamFnApi.SampleDataResponse.newBuilder();
     for (String pcollectionId : responseSamples.keySet()) {
@@ -124,7 +125,8 @@ public BeamFnApi.InstructionResponse.Builder handleDataSampleRequest(
    * @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) {
+  private Map<String, List<byte[]>> samplesFor(
+      List<String> descriptors, List<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
@@ -142,32 +144,11 @@ public Map<String, List<byte[]>> samplesFor(Set<String> descriptors, Set<String>
                   return;
                 }
 
-                samples.putIfAbsent(pcollectionId, Collections.EMPTY_LIST);
+                samples.putIfAbsent(pcollectionId, new ArrayList<>());
                 samples.get(pcollectionId).addAll(outputSampler.samples());
               });
         });
 
     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, 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(ImmutableSet.of(), pcollections);
-  }
 }

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

@@ -33,37 +33,36 @@
  */
 public class OutputSampler<T> {
   private final Coder<T> coder;
-  private final List<T> buffer = new ArrayList<>(maxElements);
   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;
+
   // Maximum number of elements in buffer.
-  private int maxElements = 10;
+  private final int maxElements;
 
   // Sampling rate.
-  private int sampleEveryN = 1000;
+  private final int sampleEveryN;
 
   // Total number of samples taken.
   private long numSamples = 0;
 
   // Index into the buffer of where to overwrite samples.
   private int resampleIndex = 0;
 
-  public OutputSampler(Coder<T> coder) {
-    this.coder = coder;
-  }
-
   public OutputSampler(Coder<T> coder, int maxElements, int sampleEveryN) {
-    this(coder);
+    this.coder = coder;
     this.maxElements = maxElements;
     this.sampleEveryN = sampleEveryN;
+    this.buffer = new ArrayList<>(this.maxElements);
   }
 
   /**
    * Samples every 1000th element or if it is part of the first 10 in the (local) PCollection.
    *
    * @param element the element to sample.
    */
-  public void sample(T element) {
+  public synchronized void sample(T element) {
     // Only sample the first 10 elements then after every `sampleEveryN`th element.
     numSamples += 1;
     if (numSamples > 10 && numSamples % sampleEveryN != 0) {
@@ -87,8 +86,17 @@ public void sample(T element) {
    */
   public List<byte[]> samples() {
     List<byte[]> ret = new ArrayList<>();
+
+    // Serializing can take a lot of CPU time for larger or complex elements. Copy the array here
+    // so as to not slow down the main processing hot path.
+    List<T> copiedBuffer;
+    synchronized (this) {
+      copiedBuffer = new ArrayList<>(buffer);
+      clear();
+    }
+
     ByteArrayOutputStream stream = new ByteArrayOutputStream();
-    for (T el : buffer) {
+    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
@@ -102,7 +110,6 @@ public List<byte[]> samples() {
       }
     }
 
-    clear();
     return ret;
   }