[Java] FirestoreIO: ramp up writes from a pipeline-wide start time

CHANGES.md

@@ -75,6 +75,7 @@
 * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)).
 * (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) that set defaults for all AWS service clients. The existing `--s3_*` options in `S3Options` override these defaults for the S3 client only ([#40528](https://github.com/apache/beam/issues/40528)).
 * (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)).
+* FirestoreIO batch writes now ramp up from a single pipeline-wide start time instead of each worker's first write, as DatastoreIO does. This adds a side input to the write transforms; set `--updateCompatibilityVersion` below 2.78.0 to keep the previous pipeline shape (Java) ([#27170](https://github.com/apache/beam/issues/27170)).
 
 ## Breaking Changes
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java

@@ -56,18 +56,23 @@
 import org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn;
 import org.apache.beam.sdk.io.gcp.firestore.FirestoreV1WriteFn.BatchWriteFnWithDeadLetterQueue;
 import org.apache.beam.sdk.io.gcp.firestore.FirestoreV1WriteFn.BatchWriteFnWithSummary;
+import org.apache.beam.sdk.options.StreamingOptions;
+import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.Reshuffle;
 import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.transforms.display.DisplayData;
 import org.apache.beam.sdk.transforms.display.HasDisplayData;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.PInput;
 import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
@@ -230,6 +235,29 @@ public final class FirestoreV1 {
 
   private FirestoreV1() {}
 
+  /**
+   * Creates a side input holding a single pipeline-wide instant, from which the write ramp-up
+   * budget of every worker grows, rather than each worker starting its own ramp-up from its first
+   * write. Follows the approach of {@code DatastoreV1.Mutate}.
+   *
+   * <p>Returns {@code null} when {@code updateCompatibilityVersion} requests a pipeline shape from
+   * before this side input was introduced.
+   */
+  private static @Nullable PCollectionView<Instant> createRampUpStartView(
+      PCollection<?> input, JodaClock clock) {
+    if (StreamingOptions.updateCompatibilityVersionLessThan(
+        input.getPipeline().getOptions(), "2.78.0")) {
+      return null;
+    }
+    return input
+        .getPipeline()
+        .apply("Create ramp-up start", Create.of("rampUpStart"))
+        .apply(
+            "Generate ramp-up start timestamp",
+            MapElements.into(TypeDescriptor.of(Instant.class)).via(ignored -> clock.instant()))
+        .apply("Ramp-up start view", View.asSingleton());
+  }
+
   /**
    * The class returned by this method provides the ability to create {@link PTransform PTransforms}
    * for read operations available in the Firestore V1 API provided by {@link
@@ -1485,16 +1513,20 @@ public BatchWriteWithSummary(
     @Override
     public PCollection<WriteSuccessSummary> expand(
         PCollection<com.google.firestore.v1.Write> input) {
-      return input.apply(
-          "batchWrite",
+      @Nullable PCollectionView<Instant> rampUpStartView = createRampUpStartView(input, clock);
+      ParDo.SingleOutput<com.google.firestore.v1.Write, WriteSuccessSummary> batchWrite =
           ParDo.of(
               new BatchWriteFnWithSummary(
                   clock,
                   firestoreStatefulComponentFactory,
                   rpcQosOptions,
                   CounterFactory.DEFAULT,
                   projectId,
-                  databaseId)));
+                  databaseId,
+                  rampUpStartView));
+      return input.apply(
+          "batchWrite",
+          rampUpStartView == null ? batchWrite : batchWrite.withSideInputs(rampUpStartView));
     }
 
     @Override
@@ -1618,16 +1650,20 @@ private BatchWriteWithDeadLetterQueue(
 
     @Override
     public PCollection<WriteFailure> expand(PCollection<com.google.firestore.v1.Write> input) {
-      return input.apply(
-          "batchWrite",
+      @Nullable PCollectionView<Instant> rampUpStartView = createRampUpStartView(input, clock);
+      ParDo.SingleOutput<com.google.firestore.v1.Write, WriteFailure> batchWrite =
           ParDo.of(
               new BatchWriteFnWithDeadLetterQueue(
                   clock,
                   firestoreStatefulComponentFactory,
                   rpcQosOptions,
                   CounterFactory.DEFAULT,
                   projectId,
-                  databaseId)));
+                  databaseId,
+                  rampUpStartView));
+      return input.apply(
+          "batchWrite",
+          rampUpStartView == null ? batchWrite : batchWrite.withSideInputs(rampUpStartView));
     }
 
     @Override

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java

@@ -53,6 +53,7 @@
 import org.apache.beam.sdk.util.BackOffUtils;
 import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollectionView;
 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.collect.ImmutableList;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -75,14 +76,16 @@ static final class BatchWriteFnWithSummary extends BaseBatchWriteFn<WriteSuccess
         RpcQosOptions rpcQosOptions,
         CounterFactory counterFactory,
         @Nullable String projectId,
-        @Nullable String databaseId) {
+        @Nullable String databaseId,
+        @Nullable PCollectionView<Instant> rampUpStartView) {
       super(
           clock,
           firestoreStatefulComponentFactory,
           rpcQosOptions,
           counterFactory,
           projectId,
-          databaseId);
+          databaseId,
+          rampUpStartView);
     }
 
     @Override
@@ -112,14 +115,16 @@ static final class BatchWriteFnWithDeadLetterQueue extends BaseBatchWriteFn<Writ
         RpcQosOptions rpcQosOptions,
         CounterFactory counterFactory,
         @Nullable String projectId,
-        @Nullable String databaseId) {
+        @Nullable String databaseId,
+        @Nullable PCollectionView<Instant> rampUpStartView) {
       super(
           clock,
           firestoreStatefulComponentFactory,
           rpcQosOptions,
           counterFactory,
           projectId,
-          databaseId);
+          databaseId,
+          rampUpStartView);
     }
 
     @Override
@@ -175,6 +180,7 @@ abstract static class BaseBatchWriteFn<OutT> extends ExplicitlyWindowedFirestore
     private transient DatabaseRootName databaseRootName;
     private final @Nullable String configuredProjectId;
     private final @Nullable String configuredDatabaseId;
+    private final @Nullable PCollectionView<Instant> rampUpStartView;
 
     @VisibleForTesting
     transient Queue<@NonNull WriteElement> writes = new PriorityQueue<>(WriteElement.COMPARATOR);
@@ -190,14 +196,16 @@ abstract static class BaseBatchWriteFn<OutT> extends ExplicitlyWindowedFirestore
         RpcQosOptions rpcQosOptions,
         CounterFactory counterFactory,
         @Nullable String configuredProjectId,
-        @Nullable String configuredDatabaseId) {
+        @Nullable String configuredDatabaseId,
+        @Nullable PCollectionView<Instant> rampUpStartView) {
       this.clock = clock;
       this.firestoreStatefulComponentFactory = firestoreStatefulComponentFactory;
       this.rpcQosOptions = rpcQosOptions;
       this.counterFactory = counterFactory;
       this.rpcAttemptContext = V1FnRpcAttemptContext.BatchWrite;
       this.configuredProjectId = configuredProjectId;
       this.configuredDatabaseId = configuredDatabaseId;
+      this.rampUpStartView = rampUpStartView;
     }
 
     @Override
@@ -265,6 +273,9 @@ public void processElement(ProcessContext context, BoundedWindow window) throws
           "nullness") // error checker is configured to treat any method not explicitly annotated as
       // @Nullable as non-null, this includes Objects.requireNonNull
       Write write = requireNonNull(context.element(), "context.element() must be non null");
+      if (rampUpStartView != null) {
+        rpcQos.setRampUpStart(context.sideInput(rampUpStartView));
+      }
       ProcessContextAdapter<OutT> contextAdapter = new ProcessContextAdapter<>(context);
       int serializedSize = write.getSerializedSize();
       boolean tooLarge = rpcQos.bytesOverLimit(serializedSize);

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQos.java

@@ -98,6 +98,15 @@ interface RpcQos {
    */
   RpcWriteAttempt newWriteAttempt(Context context);
 
+  /**
+   * Set the instant from which the write ramp-up budget grows.
+   *
+   * <p>If never called, the ramp-up starts from the first write attempted by this instance.
+   *
+   * @param rampUpStart The instant the ramp-up is considered to have started.
+   */
+  void setRampUpStart(Instant rampUpStart);
+
   /**
    * Check if a request is over the max allowed number of bytes.
    *

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosImpl.java

@@ -125,6 +125,11 @@ final class RpcQosImpl implements RpcQos {
     computeCounters = (Context c) -> O11y.create(c, counterFactory, filteringDistributionFactory);
   }
 
+  @Override
+  public void setRampUpStart(Instant rampUpStart) {
+    writeRampUp.setFirstInstant(rampUpStart);
+  }
+
   @Override
   public RpcWriteAttemptImpl newWriteAttempt(Context context) {
     return new RpcWriteAttemptImpl(
@@ -627,6 +632,10 @@ static final class WriteRampUp {
           distributionFactory.get(RpcQos.class.getName(), "qos_rampUp_availableWriteCountBudget");
     }
 
+    void setFirstInstant(Instant instant) {
+      firstInstant = Optional.of(instant);
+    }
+
     int getAvailableWriteCountBudget(Instant instant) {
       if (!firstInstant.isPresent()) {
         firstInstant = Optional.of(instant);

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1BatchWriteRampUpStartTest.java

@@ -0,0 +1,85 @@
+/*
+ * 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.io.gcp.firestore;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.firestore.v1.Write;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.extensions.protobuf.ProtoCoder;
+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.runners.TransformHierarchy;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class FirestoreV1BatchWriteRampUpStartTest {
+
+  @Test
+  public void batchWrite_usesPipelineWideRampUpStart() {
+    assertEquals(1, batchWriteSideInputs(PipelineOptionsFactory.create(), false).size());
+  }
+
+  @Test
+  public void batchWriteWithDeadLetterQueue_usesPipelineWideRampUpStart() {
+    assertEquals(1, batchWriteSideInputs(PipelineOptionsFactory.create(), true).size());
+  }
+
+  @Test
+  public void batchWrite_olderUpdateCompatibilityVersion_keepsPreviousPipelineShape() {
+    StreamingOptions options = PipelineOptionsFactory.as(StreamingOptions.class);
+    options.setUpdateCompatibilityVersion("2.77.0");
+    assertEquals(0, batchWriteSideInputs(options, false).size());
+  }
+
+  private static List<PCollectionView<?>> batchWriteSideInputs(
+      PipelineOptions options, boolean withDeadLetterQueue) {
+    Pipeline pipeline = Pipeline.create(options);
+    PCollection<Write> writes = pipeline.apply(Create.empty(ProtoCoder.of(Write.class)));
+    if (withDeadLetterQueue) {
+      writes.apply(FirestoreIO.v1().write().batchWrite().withDeadLetterQueue().build());
+    } else {
+      writes.apply(FirestoreIO.v1().write().batchWrite().build());
+    }
+
+    List<PCollectionView<?>> sideInputs = new ArrayList<>();
+    pipeline.traverseTopologically(
+        new Pipeline.PipelineVisitor.Defaults() {
+          @Override
+          public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
+            PTransform<?, ?> transform = node.getTransform();
+            if (transform instanceof ParDo.SingleOutput
+                && node.getFullName().endsWith("/batchWrite")) {
+              sideInputs.addAll(((ParDo.SingleOutput<?, ?>) transform).getSideInputs().values());
+            }
+            return CompositeBehavior.ENTER_TRANSFORM;
+          }
+        });
+    return sideInputs;
+  }
+}

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithDeadLetterQueueTest.java

@@ -223,6 +223,6 @@ protected BatchWriteFnWithDeadLetterQueue getFn(
       CounterFactory counterFactory,
       DistributionFactory distributionFactory) {
     return new BatchWriteFnWithDeadLetterQueue(
-        clock, ff, rpcQosOptions, counterFactory, null, null);
+        clock, ff, rpcQosOptions, counterFactory, null, null, null);
   }
 }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithSummaryTest.java

@@ -19,9 +19,11 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -52,6 +54,7 @@
 import org.apache.beam.sdk.io.gcp.firestore.FirestoreV1WriteFn.BatchWriteFnWithSummary;
 import org.apache.beam.sdk.io.gcp.firestore.FirestoreV1WriteFn.WriteElement;
 import org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt.Element;
+import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.joda.time.Instant;
 import org.junit.After;
@@ -201,7 +204,7 @@ public void nonRetryableWriteResultStopsAttempts() throws Exception {
 
     BaseBatchWriteFn<WriteSuccessSummary> fn =
         new BatchWriteFnWithSummary(
-            clock, ff, options, CounterFactory.DEFAULT, "testing-project", "(default)");
+            clock, ff, options, CounterFactory.DEFAULT, "testing-project", "(default)", null);
     fn.setup();
     fn.startBundle(startBundleContext);
     fn.processElement(processContext, window); // write0
@@ -238,6 +241,31 @@ public void nonRetryableWriteResultStopsAttempts() throws Exception {
     verifyNoMoreInteractions(callable);
   }
 
+  @Test
+  public void rampUpStartIsReadFromSideInput() throws Exception {
+    @SuppressWarnings("unchecked")
+    PCollectionView<Instant> rampUpStartView = mock(PCollectionView.class);
+    Instant rampUpStart = Instant.ofEpochSecond(42);
+    when(processContext.sideInput(rampUpStartView)).thenReturn(rampUpStart);
+    when(processContext.element()).thenReturn(FirestoreProtoHelpers.newWrite());
+    InterruptedException interruptedException = new InterruptedException();
+    when(attempt.awaitSafeToProceed(any())).thenThrow(interruptedException);
+
+    BaseBatchWriteFn<WriteSuccessSummary> fn =
+        new BatchWriteFnWithSummary(
+            clock, ff, rpcQosOptions, CounterFactory.DEFAULT, null, null, rampUpStartView);
+    fn.setup();
+    fn.startBundle(startBundleContext);
+    try {
+      fn.processElement(processContext, window);
+      fail("expected the attempt to be interrupted");
+    } catch (InterruptedException e) {
+      assertSame(interruptedException, e);
+    }
+
+    verify(rpcQos).setRampUpStart(rampUpStart);
+  }
+
   @Test
   public void testWithProjectId_thenWithDatabaseId() {
     FirestoreV1.Write beamWrite =
@@ -254,6 +282,6 @@ protected BatchWriteFnWithSummary getFn(
       RpcQosOptions rpcQosOptions,
       CounterFactory counterFactory,
       DistributionFactory distributionFactory) {
-    return new BatchWriteFnWithSummary(clock, ff, rpcQosOptions, counterFactory, null, null);
+    return new BatchWriteFnWithSummary(clock, ff, rpcQosOptions, counterFactory, null, null, null);
   }
 }