getProcessingTimesByStep returns a modifiable map (#29786) (#29799)

damccorm · clmccart · web-flow · commit ea477dd73b34 · 2023-12-18T11:21:42.000-05:00
* getProcessingTimesByStep returns a modifiable map

* return copy in getProcessingTimesByStep and update name accordingly

* Spotless

---------

Co-authored-by: clmccart &lt;clairelmccarthy14@gmail.com&gt;
Co-authored-by: Claire McCarthy &lt;clairemccarthy@google.com&gt;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionContext.java
@@ -341,8 +341,9 @@ public Optional<ActiveMessageMetadata> getActiveMessageMetadata() {
       return Optional.ofNullable(activeMessageMetadata);
     }
 
-    public Map<String, IntSummaryStatistics> getProcessingTimesByStep() {
-      return Collections.unmodifiableMap(processingTimesByStep);
+    public Map<String, IntSummaryStatistics> getProcessingTimesByStepCopy() {
+      Map<String, IntSummaryStatistics> processingTimesCopy = processingTimesByStep;
+      return processingTimesCopy;
     }
 
     /**
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionStateSampler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionStateSampler.java
@@ -86,7 +86,8 @@ public void removeTracker(ExecutionStateTracker tracker) {
       return;
     }
     DataflowExecutionStateTracker dfTracker = (DataflowExecutionStateTracker) tracker;
-    completedProcessingMetrics.put(dfTracker.getWorkItemId(), dfTracker.getProcessingTimesByStep());
+    completedProcessingMetrics.put(
+        dfTracker.getWorkItemId(), dfTracker.getProcessingTimesByStepCopy());
     activeTrackersByWorkId.remove(dfTracker.getWorkItemId());
 
     // Attribute any remaining time since the last sampling while removing the tracker.
@@ -126,7 +127,7 @@ public Map<String, IntSummaryStatistics> getProcessingDistributionsForWorkId(Str
     DataflowExecutionStateTracker tracker = activeTrackersByWorkId.get(workId);
     return mergeStepStatsMaps(
         completedProcessingMetrics.getOrDefault(workId, new HashMap<>()),
-        tracker.getProcessingTimesByStep());
+        tracker.getProcessingTimesByStepCopy());
   }
 
   public void resetForWorkId(String workId) {
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionContextTest.java
@@ -176,7 +176,7 @@ public void testDataflowExecutionStateTrackerRecordsCompletedProcessingTimes()
     tracker.enterState(newState);
 
     // The first completed state should be recorded and the new state should be active.
-    Map<String, IntSummaryStatistics> gotProcessingTimes = tracker.getProcessingTimesByStep();
+    Map<String, IntSummaryStatistics> gotProcessingTimes = tracker.getProcessingTimesByStepCopy();
     Assert.assertEquals(1, gotProcessingTimes.size());
     Assert.assertEquals(
         new HashSet<>(Arrays.asList(NameContextsForTests.nameContextForTest().userName())),
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionStateSamplerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionStateSamplerTest.java
@@ -76,7 +76,7 @@ public void testRemoveTrackerCompletedProcessingTimesGetsUpdated() {
     Map<String, IntSummaryStatistics> testCompletedProcessingTimes = new HashMap<>();
     testCompletedProcessingTimes.put("some-step", new IntSummaryStatistics());
     DataflowExecutionStateTracker trackerMock = createMockTracker(workId);
-    when(trackerMock.getProcessingTimesByStep()).thenReturn(testCompletedProcessingTimes);
+    when(trackerMock.getProcessingTimesByStepCopy()).thenReturn(testCompletedProcessingTimes);
 
     sampler.addTracker(trackerMock);
     sampler.removeTracker(trackerMock);
@@ -98,7 +98,7 @@ public void testGetCompletedProcessingTimesAndActiveMessageFromActiveTracker() {
         ActiveMessageMetadata.create(step1act1.getStepName().userName(), clock.getMillis());
     DataflowExecutionStateTracker trackerMock = createMockTracker(workId);
     when(trackerMock.getActiveMessageMetadata()).thenReturn(Optional.of(testMetadata));
-    when(trackerMock.getProcessingTimesByStep()).thenReturn(testCompletedProcessingTimes);
+    when(trackerMock.getProcessingTimesByStepCopy()).thenReturn(testCompletedProcessingTimes);
 
     sampler.addTracker(trackerMock);
 
@@ -122,21 +122,21 @@ public void testResetForWorkIdClearsMaps() {
         equalTo(tracker1Mock.getActiveMessageMetadata()));
     assertThat(
         sampler.getProcessingDistributionsForWorkId(workId1),
-        equalTo(tracker1Mock.getProcessingTimesByStep()));
+        equalTo(tracker1Mock.getProcessingTimesByStepCopy()));
     assertThat(
         sampler.getActiveMessageMetadataForWorkId(workId2),
         equalTo(tracker2Mock.getActiveMessageMetadata()));
     assertThat(
         sampler.getProcessingDistributionsForWorkId(workId2),
-        equalTo(tracker2Mock.getProcessingTimesByStep()));
+        equalTo(tracker2Mock.getProcessingTimesByStepCopy()));
 
     sampler.removeTracker(tracker1Mock);
     sampler.removeTracker(tracker2Mock);
     sampler.resetForWorkId(workId2);
 
     assertThat(
         sampler.getProcessingDistributionsForWorkId(workId1),
-        equalTo(tracker1Mock.getProcessingTimesByStep()));
+        equalTo(tracker1Mock.getProcessingTimesByStepCopy()));
     Assert.assertTrue(sampler.getProcessingDistributionsForWorkId(workId2).isEmpty());
   }
 

Original file line number	Diff line number	Diff line change
`@@ -341,8 +341,9 @@ public Optional<ActiveMessageMetadata> getActiveMessageMetadata() {`
`341`	`341`	`return Optional.ofNullable(activeMessageMetadata);`
`342`	`342`	`}`
`343`	`343`
`344`		`- public Map<String, IntSummaryStatistics> getProcessingTimesByStep() {`
`345`		`- return Collections.unmodifiableMap(processingTimesByStep);`
	`344`	`+ public Map<String, IntSummaryStatistics> getProcessingTimesByStepCopy() {`
	`345`	`+ Map<String, IntSummaryStatistics> processingTimesCopy = processingTimesByStep;`
	`346`	`+ return processingTimesCopy;`
`346`	`347`	`}`
`347`	`348`
`348`	`349`	`/**`