Merge branch 'master' into save-main-session-registries

Author: Abacn

.github/trigger_files/beam_PostCommit_Python.json

@@ -1,6 +1,6 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run.",
   "pr": "36271",
-  "modification": 34
+  "modification": 35
 }
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java

@@ -264,11 +264,15 @@ public long add(WindowedValue<T> data) throws IOException {
         }
       }
 
-      Windmill.KeyedMessageBundle.Builder keyedOutput = productionMap.get(key);
-      if (keyedOutput == null) {
-        keyedOutput = Windmill.KeyedMessageBundle.newBuilder().setKey(key);
-        productionMap.put(key, keyedOutput);
-      }
+      Windmill.KeyedMessageBundle.Builder keyedOutput =
+          productionMap.computeIfAbsent(
+              key,
+              (k) -> {
+                Windmill.KeyedMessageBundle.Builder builder =
+                    Windmill.KeyedMessageBundle.newBuilder();
+                builder.setKey(k);
+                return builder;
+              });
 
       try {
         messageBuilder

sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/WeightedList.java

@@ -20,6 +20,7 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicLong;
 import org.apache.beam.sdk.util.Weighted;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.math.LongMath;
 
 /** Facade for a {@link List<T>} that keeps track of weight, for cache limit reasons. */
 public class WeightedList<T> implements Weighted {
@@ -71,14 +72,6 @@ public void addAll(List<T> values, long weight) {
   }
 
   public void accumulateWeight(long weight) {
-    this.weight.accumulateAndGet(
-        weight,
-        (first, second) -> {
-          try {
-            return Math.addExact(first, second);
-          } catch (ArithmeticException e) {
-            return Long.MAX_VALUE;
-          }
-        });
+    this.weight.accumulateAndGet(weight, LongMath::saturatedAdd);
   }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackersTest.java

@@ -102,7 +102,7 @@ private static class RestrictionTrackerWithProgress extends RestrictionTracker<O
       implements HasProgress {
     private boolean blockTryClaim;
     private boolean blockTrySplit;
-    private boolean isBlocked;
+    private volatile boolean isBlocked;
     public static final Progress REPORT_PROGRESS = Progress.from(2.0, 3.0);
 
     public RestrictionTrackerWithProgress() {

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java

@@ -49,6 +49,8 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.math.LongMath;
 
 /**
  * Adapters which convert a logical series of chunks using continuation tokens over the Beam Fn
@@ -249,15 +251,11 @@ static class BlocksPrefix<T> extends Blocks<T> implements Shrinkable<BlocksPrefi
 
       @Override
       public long getWeight() {
-        try {
-          long sum = 8 + blocks.size() * 8L;
-          for (Block<T> block : blocks) {
-            sum = Math.addExact(sum, block.getWeight());
-          }
-          return sum;
-        } catch (ArithmeticException e) {
-          return Long.MAX_VALUE;
+        long sum = 8 + blocks.size() * 8L;
+        for (Block<T> block : blocks) {
+          sum = LongMath.saturatedAdd(sum, block.getWeight());
         }
+        return sum;
       }
 
       BlocksPrefix(List<Block<T>> blocks) {
@@ -282,8 +280,7 @@ public List<Block<T>> getBlocks() {
 
     @AutoValue
     abstract static class Block<T> implements Weighted {
-      private static final Block<Void> EMPTY =
-          fromValues(WeightedList.of(Collections.emptyList(), 0), null);
+      private static final Block<Void> EMPTY = fromValues(ImmutableList.of(), 0, null);
 
       @SuppressWarnings("unchecked") // Based upon as Collections.emptyList()
       public static <T> Block<T> emptyBlock() {
@@ -299,21 +296,37 @@ public static <T> Block<T> mutatedBlock(WeightedList<T> values) {
       }
 
       public static <T> Block<T> fromValues(List<T> values, @Nullable ByteString nextToken) {
-        return fromValues(WeightedList.of(values, Caches.weigh(values)), nextToken);
+        if (values.isEmpty() && nextToken == null) {
+          return emptyBlock();
+        }
+        ImmutableList<T> immutableValues = ImmutableList.copyOf(values);
+        long listWeight = immutableValues.size() * Caches.REFERENCE_SIZE;
+        for (T value : immutableValues) {
+          listWeight = LongMath.saturatedAdd(listWeight, Caches.weigh(value));
+        }
+        return fromValues(immutableValues, listWeight, nextToken);
       }
 
       public static <T> Block<T> fromValues(
           WeightedList<T> values, @Nullable ByteString nextToken) {
-        long weight = values.getWeight() + 24;
+        if (values.isEmpty() && nextToken == null) {
+          return emptyBlock();
+        }
+        return fromValues(ImmutableList.copyOf(values.getBacking()), values.getWeight(), nextToken);
+      }
+
+      private static <T> Block<T> fromValues(
+          ImmutableList<T> values, long listWeight, @Nullable ByteString nextToken) {
+        long weight = LongMath.saturatedAdd(listWeight, 24);
         if (nextToken != null) {
           if (nextToken.isEmpty()) {
             nextToken = ByteString.EMPTY;
           } else {
-            weight += Caches.weigh(nextToken);
+            weight = LongMath.saturatedAdd(weight, Caches.weigh(nextToken));
           }
         }
         return new AutoValue_StateFetchingIterators_CachingStateIterable_Block<>(
-            values.getBacking(), nextToken, weight);
+            values, nextToken, weight);
       }
 
       abstract List<T> getValues();
@@ -372,10 +385,12 @@ public void remove(Set<Object> toRemoveStructuralValues) {
         totalSize += tBlock.getValues().size();
       }
 
-      WeightedList<T> allValues = WeightedList.of(new ArrayList<>(totalSize), 0L);
+      ImmutableList.Builder<T> allValues = ImmutableList.builderWithExpectedSize(totalSize);
+      long weight = 0;
+      List<T> blockValuesToKeep = new ArrayList<>();
       for (Block<T> block : blocks) {
+        blockValuesToKeep.clear();
         boolean valueRemovedFromBlock = false;
-        List<T> blockValuesToKeep = new ArrayList<>();
         for (T value : block.getValues()) {
           if (!toRemoveStructuralValues.contains(valueCoder.structuralValue(value))) {
             blockValuesToKeep.add(value);
@@ -387,13 +402,19 @@ public void remove(Set<Object> toRemoveStructuralValues) {
         // If any value was removed from this block, need to estimate the weight again.
         // Otherwise, just reuse the block's weight.
         if (valueRemovedFromBlock) {
-          allValues.addAll(blockValuesToKeep, Caches.weigh(block.getValues()));
+          allValues.addAll(blockValuesToKeep);
+          for (T value : blockValuesToKeep) {
+            weight = LongMath.saturatedAdd(weight, Caches.weigh(value));
+          }
         } else {
-          allValues.addAll(block.getValues(), block.getWeight());
+          allValues.addAll(block.getValues());
+          weight = LongMath.saturatedAdd(weight, block.getWeight());
         }
       }
 
-      cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues)));
+      cache.put(
+          IterableCacheKey.INSTANCE,
+          new MutatedBlocks<>(Block.fromValues(allValues.build(), weight, null)));
     }
 
     /**
@@ -484,21 +505,24 @@ private void appendHelper(List<T> newValues, long newWeight) {
       for (Block<T> block : blocks) {
         totalSize += block.getValues().size();
       }
-      WeightedList<T> allValues = WeightedList.of(new ArrayList<>(totalSize), 0L);
+      ImmutableList.Builder<T> allValues = ImmutableList.builderWithExpectedSize(totalSize);
+      long weight = 0;
       for (Block<T> block : blocks) {
-        allValues.addAll(block.getValues(), block.getWeight());
+        allValues.addAll(block.getValues());
+        weight = LongMath.saturatedAdd(weight, block.getWeight());
       }
       if (newWeight < 0) {
-        if (newValues.size() == 1) {
-          // Optimize weighing of the common value state as single single-element bag state.
-          newWeight = Caches.weigh(newValues.get(0));
-        } else {
-          newWeight = Caches.weigh(newValues);
+        newWeight = 0;
+        for (T value : newValues) {
+          newWeight = LongMath.saturatedAdd(newWeight, Caches.weigh(value));
         }
       }
-      allValues.addAll(newValues, newWeight);
+      allValues.addAll(newValues);
+      weight = LongMath.saturatedAdd(weight, newWeight);
 
-      cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues)));
+      cache.put(
+          IterableCacheKey.INSTANCE,
+          new MutatedBlocks<>(Block.fromValues(allValues.build(), weight, null)));
     }
 
     class CachingStateIterator implements PrefetchableIterator<T> {
@@ -580,8 +604,7 @@ public boolean hasNext() {
             return false;
           }
           // Release the block while we are loading the next one.
-          currentBlock =
-              Block.fromValues(WeightedList.of(Collections.emptyList(), 0L), ByteString.EMPTY);
+          currentBlock = Block.emptyBlock();
 
           @Nullable Blocks<T> existing = cache.peek(IterableCacheKey.INSTANCE);
           boolean isFirstBlock = ByteString.EMPTY.equals(nextToken);

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java

@@ -70,6 +70,7 @@
 import org.joda.time.Duration;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -115,6 +116,7 @@ public void tearDown() throws NoSuchFieldException, IllegalAccessException {
   }
 
   @Test
+  @Ignore("https://github.com/apache/beam/issues/37002 Re-enable skipped tests.")
   // Error code UNAVAILABLE is retried repeatedly until the RPC times out.
   public void testUnavailableExceptionRetries() throws InterruptedException {
     DirectOptions options = PipelineOptionsFactory.as(DirectOptions.class);
@@ -155,6 +157,7 @@ public void testUnavailableExceptionRetries() throws InterruptedException {
   }
 
   @Test
+  @Ignore("https://github.com/apache/beam/issues/37002 Re-enable skipped tests.")
   // Error code ABORTED is retried repeatedly until it times out.
   public void testAbortedExceptionRetries() throws InterruptedException {
     mockSpannerService.setExecuteStreamingSqlExecutionTime(
@@ -218,6 +221,7 @@ public void testUnknownExceptionDoesNotRetry() {
   }
 
   @Test
+  @Ignore("https://github.com/apache/beam/issues/37002 Re-enable skipped tests.")
   // Error code RESOURCE_EXHAUSTED is retried repeatedly.
   public void testResourceExhaustedRetry() {
     mockSpannerService.setExecuteStreamingSqlExecutionTime(
@@ -281,6 +285,7 @@ public void testResourceExhaustedRetryWithDefaultSettings() {
   }
 
   @Test
+  @Ignore("https://github.com/apache/beam/issues/37002 Re-enable skipped tests.")
   public void testInvalidRecordReceived() {
     final Timestamp startTimestamp = Timestamp.ofTimeSecondsAndNanos(0, 1000);
     final Timestamp endTimestamp =

sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java

@@ -25,6 +25,7 @@
 import static org.apache.beam.sdk.io.jms.CommonJms.toSerializableFunction;
 import static org.apache.beam.sdk.io.jms.JmsIO.Writer.JMS_IO_PRODUCER_METRIC_NAME;
 import static org.apache.beam.sdk.io.jms.JmsIO.Writer.PUBLICATION_RETRIES_METRIC_NAME;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
 import static org.hamcrest.CoreMatchers.allOf;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.contains;
@@ -86,6 +87,7 @@
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
 import org.apache.beam.sdk.io.jms.JmsIO.UnboundedJmsReader;
 import org.apache.beam.sdk.metrics.MetricNameFilter;
@@ -541,6 +543,16 @@ public void testSplitForTopic() throws Exception {
     assertEquals(1, splits.size());
   }
 
+  private boolean advanceWithRetry(UnboundedSource.UnboundedReader reader) throws IOException {
+    for (int attempt = 0; attempt < 10; attempt++) {
+      if (reader.advance()) {
+        return true;
+      }
+      sleepUninterruptibly(java.time.Duration.ofMillis(100));
+    }
+    return false;
+  }
+
   @Test
   public void testCheckpointMark() throws Exception {
     // we are using no prefetch here
@@ -558,7 +570,7 @@ public void testCheckpointMark() throws Exception {
 
     // consume 3 messages (NB: start already consumed the first message)
     for (int i = 0; i < 3; i++) {
-      assertTrue(String.format("Failed at %d-th message", i), reader.advance());
+      assertTrue(String.format("Failed at %d-th message", i), advanceWithRetry(reader));
     }
 
     // the messages are still pending in the queue (no ACK yet)
@@ -572,7 +584,7 @@ public void testCheckpointMark() throws Exception {
 
     // we read the 6 pending messages
     for (int i = 0; i < 6; i++) {
-      assertTrue(String.format("Failed at %d-th message", i), reader.advance());
+      assertTrue(String.format("Failed at %d-th message", i), advanceWithRetry(reader));
     }
 
     // still 6 pending messages as we didn't finalize the checkpoint
@@ -592,8 +604,8 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception {
     assertTrue(reader.start());
 
     // consume 2 message (NB: start already consumed the first message)
-    assertTrue(reader.advance());
-    assertTrue(reader.advance());
+    assertTrue(advanceWithRetry(reader));
+    assertTrue(advanceWithRetry(reader));
 
     // get checkpoint mark after consumed 4 messages
     CheckpointMark mark = reader.getCheckpointMark();
@@ -724,7 +736,7 @@ public void testCheckpointMarkSafety() throws Exception {
 
     // consume half the messages (NB: start already consumed the first message)
     for (int i = 0; i < (messagesToProcess / 2) - 1; i++) {
-      assertTrue(reader.advance());
+      assertTrue(advanceWithRetry(reader));
     }
 
     // the messages are still pending in the queue (no ACK yet)
@@ -738,7 +750,7 @@ public void testCheckpointMarkSafety() throws Exception {
             () -> {
               try {
                 for (int i = 0; i < messagesToProcess / 2; i++) {
-                  assertTrue(reader.advance());
+                  assertTrue(advanceWithRetry(reader));
                 }
               } catch (IOException ex) {
                 throw new RuntimeException(ex);
@@ -877,7 +889,7 @@ public void testDiscardCheckpointMark() throws Exception {
 
     // consume 3 more messages (NB: start already consumed the first message)
     for (int i = 0; i < 3; i++) {
-      assertTrue(reader.advance());
+      assertTrue(advanceWithRetry(reader));
     }
 
     // the messages are still pending in the queue (no ACK yet)
@@ -891,7 +903,7 @@ public void testDiscardCheckpointMark() throws Exception {
 
     // we read the 6 pending messages
     for (int i = 0; i < 6; i++) {
-      assertTrue(reader.advance());
+      assertTrue(advanceWithRetry(reader));
     }
 
     // still 6 pending messages as we didn't finalize the checkpoint