Fixing State.readLater() usage + Unit test fix to reflect changed exp…

…ectations
apache · Nov 16, 2022 · 1e1770b · 1e1770b
1 parent 5840e55
commit 1e1770b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
@@ -464,8 +464,21 @@ public void processElement(
         @Timestamp Instant elementTs,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
+
+      final boolean shouldCareAboutWeight = weigher != null && batchSizeBytes != Long.MAX_VALUE;
+      final boolean shouldCareAboutMaxBufferingDuration =
+          maxBufferingDuration.isLongerThan(Duration.ZERO);
+
+      if (shouldCareAboutWeight) {
+        storedBatchSizeBytes.readLater();
+      }
+      storedBatchSize.readLater();
+      if (shouldCareAboutMaxBufferingDuration) {
+        minBufferedTs.readLater();
+      }
+
       LOG.debug("*** BATCH *** Add element for window {} ", window);
-      if (shouldCareAboutWeight()) {
+      if (shouldCareAboutWeight) {
         final long elementWeight = weigher.apply(element.getValue());
         if (elementWeight + storedBatchSizeBytes.read() > batchSizeBytes) {
           // Firing by count and size limits behave differently.
@@ -494,17 +507,13 @@ public void processElement(
           bufferingTimer.clear();
         }
         storedBatchSizeBytes.add(elementWeight);
-        storedBatchSizeBytes.readLater();
       }
       batch.add(element.getValue());
       // Blind add is supported with combiningState
       storedBatchSize.add(1L);
 
-      long num;
-      if (maxBufferingDuration.isLongerThan(Duration.ZERO)) {
-        minBufferedTs.readLater();
-        num = storedBatchSize.read();
-
+      final long num = storedBatchSize.read();
+      if (shouldCareAboutMaxBufferingDuration) {
         long oldOutputTs =
             MoreObjects.firstNonNull(
                 minBufferedTs.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
@@ -523,15 +532,14 @@ public void processElement(
               .set(Instant.ofEpochMilli(targetTs));
         }
       }
-      num = storedBatchSize.read();
 
       if (num % prefetchFrequency == 0) {
         // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
 
       if (num >= batchSize
-          || (shouldCareAboutWeight() && storedBatchSizeBytes.read() >= batchSizeBytes)) {
+          || (shouldCareAboutWeight && storedBatchSizeBytes.read() >= batchSizeBytes)) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
         flushBatch(
             receiver,
@@ -545,10 +553,6 @@ public void processElement(
       }
     }
 
-    private boolean shouldCareAboutWeight() {
-      return weigher != null && batchSizeBytes != Long.MAX_VALUE;
-    }
-
     @OnTimer(END_OF_BUFFERING_ID)
     public void onBufferingTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,

diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java
@@ -158,7 +158,7 @@ public Void apply(Iterable<KV<String, Iterable<String>>> input) {
               }
             });
     PAssert.thatSingleton("Incorrect collection size", collection.apply("Count", Count.globally()))
-        .isEqualTo(3L);
+        .isEqualTo(4L);
     pipeline.run();
   }
 
@@ -196,7 +196,7 @@ public Void apply(Iterable<KV<String, Iterable<String>>> input) {
               }
             });
     PAssert.thatSingleton("Incorrect collection size", collection.apply("Count", Count.globally()))
-        .isEqualTo(5L);
+        .isEqualTo(9L);
     pipeline.run();
   }