[Cherry Pick #29979] Fix race condition in Dataflow sampler. (#30144)

* make getProcessingDistributionsForWorkId synchronized to avoid race conditions with remove tracker * correctly access the concurrent hashmap to avoid race conditions --------- Co-authored-by: Claire McCarthy <[email protected]>
apache · Jan 30, 2024 · 1c974aa · 1c974aa
1 parent d1d3457
commit 1c974aa
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/.../src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionStateSampler.java b/.../src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionStateSampler.java
@@ -118,13 +118,10 @@ public Optional<ActiveMessageMetadata> getActiveMessageMetadataForWorkId(String
   }
 
   public Map<String, IntSummaryStatistics> getProcessingDistributionsForWorkId(String workId) {
-    if (!activeTrackersByWorkId.containsKey(workId)) {
-      if (completedProcessingMetrics.containsKey(workId)) {
-        return completedProcessingMetrics.get(workId);
-      }
-      return new HashMap<>();
-    }
     DataflowExecutionStateTracker tracker = activeTrackersByWorkId.get(workId);
+    if (tracker == null) {
+      return completedProcessingMetrics.getOrDefault(workId, new HashMap<>());
+    }
     return mergeStepStatsMaps(
         completedProcessingMetrics.getOrDefault(workId, new HashMap<>()),
         tracker.getProcessingTimesByStepCopy());