diff --git a/CHANGES.md b/CHANGES.md index 7693e65d78a0..c3266e5bf24a 100644 --- a/CHANGES.md +++ b/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 diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java index e8da6097788a..552ff3e9011e 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java +++ b/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}. + * + *

Returns {@code null} when {@code updateCompatibilityVersion} requests a pipeline shape from + * before this side input was introduced. + */ + private static @Nullable PCollectionView 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,8 +1513,8 @@ public BatchWriteWithSummary( @Override public PCollection expand( PCollection input) { - return input.apply( - "batchWrite", + @Nullable PCollectionView rampUpStartView = createRampUpStartView(input, clock); + ParDo.SingleOutput batchWrite = ParDo.of( new BatchWriteFnWithSummary( clock, @@ -1494,7 +1522,11 @@ public PCollection expand( rpcQosOptions, CounterFactory.DEFAULT, projectId, - databaseId))); + databaseId, + rampUpStartView)); + return input.apply( + "batchWrite", + rampUpStartView == null ? batchWrite : batchWrite.withSideInputs(rampUpStartView)); } @Override @@ -1618,8 +1650,8 @@ private BatchWriteWithDeadLetterQueue( @Override public PCollection expand(PCollection input) { - return input.apply( - "batchWrite", + @Nullable PCollectionView rampUpStartView = createRampUpStartView(input, clock); + ParDo.SingleOutput batchWrite = ParDo.of( new BatchWriteFnWithDeadLetterQueue( clock, @@ -1627,7 +1659,11 @@ public PCollection expand(PCollection rampUpStartView) { super( clock, firestoreStatefulComponentFactory, rpcQosOptions, counterFactory, projectId, - databaseId); + databaseId, + rampUpStartView); } @Override @@ -112,14 +115,16 @@ static final class BatchWriteFnWithDeadLetterQueue extends BaseBatchWriteFn rampUpStartView) { super( clock, firestoreStatefulComponentFactory, rpcQosOptions, counterFactory, projectId, - databaseId); + databaseId, + rampUpStartView); } @Override @@ -175,6 +180,7 @@ abstract static class BaseBatchWriteFn extends ExplicitlyWindowedFirestore private transient DatabaseRootName databaseRootName; private final @Nullable String configuredProjectId; private final @Nullable String configuredDatabaseId; + private final @Nullable PCollectionView rampUpStartView; @VisibleForTesting transient Queue<@NonNull WriteElement> writes = new PriorityQueue<>(WriteElement.COMPARATOR); @@ -190,7 +196,8 @@ abstract static class BaseBatchWriteFn extends ExplicitlyWindowedFirestore RpcQosOptions rpcQosOptions, CounterFactory counterFactory, @Nullable String configuredProjectId, - @Nullable String configuredDatabaseId) { + @Nullable String configuredDatabaseId, + @Nullable PCollectionView rampUpStartView) { this.clock = clock; this.firestoreStatefulComponentFactory = firestoreStatefulComponentFactory; this.rpcQosOptions = rpcQosOptions; @@ -198,6 +205,7 @@ abstract static class BaseBatchWriteFn extends ExplicitlyWindowedFirestore 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 contextAdapter = new ProcessContextAdapter<>(context); int serializedSize = write.getSerializedSize(); boolean tooLarge = rpcQos.bytesOverLimit(serializedSize); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQos.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQos.java index 2b0250a53af3..0b9151c61e8b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQos.java +++ b/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. + * + *

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. * diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosImpl.java index 06b7dbe38013..b89001ca5c16 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosImpl.java +++ b/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); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1BatchWriteRampUpStartTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1BatchWriteRampUpStartTest.java new file mode 100644 index 000000000000..72a2cfa86cbc --- /dev/null +++ b/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> batchWriteSideInputs( + PipelineOptions options, boolean withDeadLetterQueue) { + Pipeline pipeline = Pipeline.create(options); + PCollection 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> 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; + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithDeadLetterQueueTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithDeadLetterQueueTest.java index b8f3f6ea5ddd..4cf6d42a1cd1 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithDeadLetterQueueTest.java +++ b/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); } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithSummaryTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithSummaryTest.java index 973521b5b055..26963769a50e 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnBatchWriteWithSummaryTest.java +++ b/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 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 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 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); } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosTest.java index 8910d491c4fd..5bbfc17345f6 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosTest.java @@ -562,6 +562,20 @@ public void initialBatchSizeRelativeToWorkerCount_1() { doTest_initialBatchSizeRelativeToWorkerCount(1, 500); } + @Test + public void rampUp_growsFromProvidedStart() { + RpcQosOptions options = + RpcQosOptions.newBuilder().withHintMaxNumWorkers(500).withBatchInitialCount(500).build(); + RpcQosImpl qos = new RpcQosImpl(options, random, sleeper, counterFactory, distributionFactory); + qos.setRampUpStart(Instant.EPOCH); + + // 20 minutes after the ramp-up start: 500 / 500 * 1.5^((20 - 5) / 5) = 3.375 + Instant firstWrite = Instant.EPOCH.plus(Duration.standardMinutes(20)); + RpcWriteAttemptImpl attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT); + FlushBufferImpl> buffer = attempt.newFlushBuffer(firstWrite); + assertEquals(3, buffer.nextBatchMaxCount); + } + @Test public void isCodeRetryable() { doTest_isCodeRetryable(Code.ABORTED, true);