Merge pull request #35714 from Abacn/reapply

Reapply "Refactor: separate SplittableTruncateSizedRestrictions"

Author: Abacn

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java

@@ -49,9 +49,11 @@
 import org.apache.beam.sdk.metrics.MetricQueryResults;
 import org.apache.beam.sdk.metrics.MetricResult;
 import org.apache.beam.sdk.metrics.MetricsFilter;
+import org.apache.beam.sdk.options.ExperimentalOptions;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.StreamingOptions;
+import org.apache.beam.sdk.util.ReleaseInfo;
 import org.apache.beam.sdk.util.common.ReflectHelpers;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
@@ -312,6 +314,13 @@ public LaunchInfo launch(String project, String region, LaunchConfig options) th
             optionFromConfig.add(
                 String.format("--tempLocation=%s", pipelineOptions.getTempLocation()));
           }
+          // for snapshot Beam running on runner v2, need to use staged SDK harness
+          if (ReleaseInfo.getReleaseInfo().isDevSdkVersion()
+              && (ExperimentalOptions.hasExperiment(pipelineOptions, "use_runner_v2")
+                  || (!Strings.isNullOrEmpty(options.parameters().get("experiments"))
+                      && options.parameters().get("experiments").contains("use_runner_v2")))) {
+            optionFromConfig.add("--experiments=use_staged_dataflow_worker_jar");
+          }
 
           // dataflow runner specific options
           PipelineOptions updatedOptions =

sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java

@@ -119,4 +119,25 @@ public static <RestrictionT, PositionT> RestrictionTracker<RestrictionT, Positio
       return new RestrictionTrackerObserver<>(restrictionTracker, claimObserver);
     }
   }
+
+  public static <RestrictionT, PositionT> RestrictionTracker<RestrictionT, PositionT> synchronize(
+      RestrictionTracker<RestrictionT, PositionT> restrictionTracker) {
+    if (restrictionTracker instanceof RestrictionTracker.HasProgress) {
+      return new RestrictionTrackerObserverWithProgress<>(
+          restrictionTracker, (ClaimObserver<PositionT>) NOOP_CLAIM_OBSERVER);
+    } else {
+      return new RestrictionTrackerObserver<>(
+          restrictionTracker, (ClaimObserver<PositionT>) NOOP_CLAIM_OBSERVER);
+    }
+  }
+
+  static class NoopClaimObserver<PositionT> implements ClaimObserver<PositionT> {
+    @Override
+    public void onClaimed(PositionT position) {}
+
+    @Override
+    public void onClaimFailed(PositionT position) {}
+  }
+
+  private static final NoopClaimObserver<Object> NOOP_CLAIM_OBSERVER = new NoopClaimObserver<>();
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java

@@ -211,6 +211,6 @@ public static <RestrictionT> TruncateResult of(RestrictionT restriction) {
       return new AutoValue_RestrictionTracker_TruncateResult(restriction);
     }
 
-    public abstract @Nullable RestrictionT getTruncatedRestriction();
+    public abstract RestrictionT getTruncatedRestriction();
   }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/util/Holder.java

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.util;
+
+import org.apache.beam.sdk.annotations.Internal;
+
+/**
+ * A trivial boxing of a value, used when nullability needs to be added to a generic type. (Optional
+ * does not work for this)
+ *
+ * <p>Example: For a generic type `T` the actual parameter may be nullable or not. So you cannot
+ * check values for null to determine presence/absence. Instead you can store a {@code @Nullable
+ * Holder<T>}.
+ */
+@Internal
+public class Holder<T> {
+  private T value;
+
+  private Holder(T value) {
+    this.value = value;
+  }
+
+  public static <ValueT> Holder<ValueT> of(ValueT value) {
+    return new Holder<>(value);
+  }
+
+  public T get() {
+    return value;
+  };
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/values/KV.java

@@ -40,6 +40,7 @@
  */
 public class KV<K, V> implements Serializable {
   /** Returns a {@link KV} with the given key and value. */
+  @Pure
   public static <K, V> KV<K, V> of(K key, V value) {
     return new KV<>(key, value);
   }

sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java

@@ -20,6 +20,7 @@
 import java.util.Collection;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.checkerframework.dataflow.qual.Pure;
 import org.joda.time.Instant;
 
 /**
@@ -29,26 +30,32 @@
  */
 public interface WindowedValue<T> {
   /** The primary data for this value. */
+  @Pure
   T getValue();
 
   /** The timestamp of this value in event time. */
+  @Pure
   Instant getTimestamp();
 
   /** Returns the windows of this {@code WindowedValue}. */
+  @Pure
   Collection<? extends BoundedWindow> getWindows();
 
   /** The {@link PaneInfo} associated with this WindowedValue. */
+  @Pure
   PaneInfo getPaneInfo();
 
   /**
    * A representation of each of the actual values represented by this compressed {@link
    * WindowedValue}, one per window.
    */
+  @Pure
   Iterable<WindowedValue<T>> explodeWindows();
 
   /**
    * A {@link WindowedValue} with identical metadata to the current one, but with the provided
    * value.
    */
+  @Pure
   <OtherT> WindowedValue<OtherT> withValue(OtherT value);
 }