Fix ConcurrentModificationException in OutputSampler

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

@@ -46,7 +46,7 @@ public class OutputSampler<T> {
   // Temporarily holds exceptional elements. These elements can also be duplicated in the main
   // buffer. This is in order to always track exceptional elements even if the number of samples in
   // the main buffer drops it.
-  private final Map<String, ElementSample<T>> exceptions = new HashMap<>();
+  private Map<String, ElementSample<T>> exceptions = new HashMap<>();
 
   // Maximum number of elements in buffer.
   private final int maxElements;
@@ -198,26 +198,30 @@ public List<BeamFnApi.SampledElement> samples() throws IOException {
     // 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<ElementSample<T>> bufferToSend;
+    Map<String, ElementSample<T>> exceptionsToSend;
     int sampleIndex = 0;
     synchronized (this) {
       bufferToSend = buffer;
-      sampleIndex = resampleIndex;
       buffer = new ArrayList<>(maxElements);
+
+      exceptionsToSend = exceptions;
+      exceptions = new HashMap<>(exceptions.size());
+
+      sampleIndex = resampleIndex;
       resampleIndex = 0;
     }
 
     // An element can live in both the main samples and exception buffer. Use a small look up table
     // to deduplicate samples.
     HashSet<Long> seen = new HashSet<>();
     ByteStringOutputStream stream = new ByteStringOutputStream();
-    for (Map.Entry<String, ElementSample<T>> pair : exceptions.entrySet()) {
+    for (Map.Entry<String, ElementSample<T>> pair : exceptionsToSend.entrySet()) {
       String processBundleId = pair.getKey();
       ElementSample<T> sample = pair.getValue();
       seen.add(sample.id);
 
       ret.add(sampleToProto(sample, stream, processBundleId));
     }
-    exceptions.clear();
 
     for (int i = 0; i < bufferToSend.size(); i++) {
       int index = (sampleIndex + i) % bufferToSend.size();

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

@@ -307,7 +307,9 @@ public void testConcurrentSamples() throws IOException, InterruptedException {
               }
 
               for (int i = 0; i < 1000000; i++) {
-                outputSampler.sample(WindowedValue.valueInGlobalWindow(i));
+                ElementSample<Integer> sample =
+                    outputSampler.sample(WindowedValue.valueInGlobalWindow(i));
+                outputSampler.exception(sample, new RuntimeException(""), "ptransformId", "pbId");
               }
 
               doneSignal.countDown();
@@ -324,7 +326,9 @@ public void testConcurrentSamples() throws IOException, InterruptedException {
               }
 
               for (int i = -1000000; i < 0; i++) {
-                outputSampler.sample(WindowedValue.valueInGlobalWindow(i));
+                ElementSample<Integer> sample =
+                    outputSampler.sample(WindowedValue.valueInGlobalWindow(i));
+                outputSampler.exception(sample, new RuntimeException(""), "ptransformId", "pbId");
               }
 
               doneSignal.countDown();