Skip to content

[SpannerIO] Add withWriteResults() to SpannerIO.Write to emit commit metadata #40480

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stankiewicz wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[SpannerIO] Add withWriteResults() to SpannerIO.Write to emit commit metadata #40480

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 140 additions & 16 deletions .../io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import static org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants.LOW_LATENCY_REAL_TIME_CHECKPOINT_INTERVAL;
import static org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants.MAX_INCLUSIVE_END_AT;
import static org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants.THROUGHPUT_WINDOW_SECONDS;
import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;

Expand All @@ -41,6 +42,7 @@
import com.google.cloud.ServiceFactory;
import com.google.cloud.Timestamp;
import com.google.cloud.spanner.AbortedException;
import com.google.cloud.spanner.CommitResponse;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Dialect;
Expand Down Expand Up @@ -529,6 +531,7 @@ public static Write write() {
.setMaxNumMutations(DEFAULT_MAX_NUM_MUTATIONS)
.setMaxNumRows(DEFAULT_MAX_NUM_ROWS)
.setFailureMode(FailureMode.FAIL_FAST)
.setCommitStats(false)
.build();
}

Expand Down Expand Up @@ -1427,6 +1430,8 @@ public abstract static class Write extends PTransform<PCollection<Mutation>, Spa

abstract @Nullable PCollectionView<Dialect> getDialectView();

abstract boolean getCommitStats();

abstract Builder toBuilder();

@AutoValue.Builder
Expand All @@ -1448,6 +1453,8 @@ abstract static class Builder {

abstract Builder setDialectView(PCollectionView<Dialect> dialect);

abstract Builder setCommitStats(boolean commitStats);

abstract Write build();
}

Expand Down Expand Up @@ -1693,6 +1700,14 @@ public Write withHighPriority() {
return withSpannerConfig(config.withRpcPriority(RpcPriority.HIGH));
}

/**
* Returns a {@link WriteWithResults} transform that enables {@link Options#commitStats()} and
* emits a {@link WriteResults} for each batch of mutations written to Cloud Spanner.
*/
public WriteWithResults withWriteResults() {
return new WriteWithResults(toBuilder().setCommitStats(true).build());
}

@Override
public SpannerWriteResult expand(PCollection<Mutation> input) {
getSpannerConfig().validate();
Expand Down Expand Up @@ -1731,6 +1746,65 @@ private void populateDisplayDataWithParamaters(DisplayData.Builder builder) {
}
}

/**
* A {@link PTransform} that writes {@link Mutation} objects to Google Cloud Spanner with {@link
* Options#commitStats()} enabled and emits a {@link WriteResults} for each committed batch.
*/
public static class WriteWithResults
extends PTransform<PCollection<Mutation>, PCollection<WriteResults>> {

private final Write spec;

WriteWithResults(Write spec) {
this.spec = spec;
}

/** Same transform but can be applied to {@link PCollection} of {@link MutationGroup}. */
public WriteGroupedWithResults grouped() {
return new WriteGroupedWithResults(spec);
}

@Override
public PCollection<WriteResults> expand(PCollection<Mutation> input) {
spec.getSpannerConfig().validate();

return input
.apply("To mutation group", ParDo.of(new ToMutationGroupFn()))
.apply("Write mutations to Cloud Spanner", new WriteGroupedWithResults(spec));
}

@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
spec.populateDisplayDataWithParamaters(builder);
}
}

/** Same as {@link WriteWithResults} but supports grouped mutations. */
public static class WriteGroupedWithResults
extends PTransform<PCollection<MutationGroup>, PCollection<WriteResults>> {

private final Write spec;

public WriteGroupedWithResults(Write spec) {
this.spec = spec;
}

@Override
public PCollection<WriteResults> expand(PCollection<MutationGroup> input) {
SpannerWriteResult result = input.apply(new WriteGrouped(spec));
return checkStateNotNull(
result.getWriteResults(),
"WriteResults PCollection must be present when commitStats is enabled.");
}

@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
spec.populateDisplayDataWithParamaters(builder);
}
}

static class WriteRows extends PTransform<PCollection<Row>, PDone> {

private final Write write;
Expand Down Expand Up @@ -1769,6 +1843,8 @@ public static class WriteGrouped
new TupleTag<Iterable<MutationGroup>>("unbatchableMutations") {};

private static final TupleTag<Void> MAIN_OUT_TAG = new TupleTag<Void>("mainOut") {};
private static final TupleTag<WriteResults> WRITE_RESULTS_TAG =
new TupleTag<WriteResults>("writeResults") {};
private static final TupleTag<MutationGroup> FAILED_MUTATIONS_TAG =
new TupleTag<MutationGroup>("failedMutations") {};
private static final SerializableCoder<MutationGroup> CODER =
Expand All @@ -1778,6 +1854,15 @@ public WriteGrouped(Write spec) {
this.spec = spec;
}

/**
* Returns a {@link WriteGroupedWithResults} transform that enables {@link
* Options#commitStats()} and emits a {@link WriteResults} for each batch of mutations written
* to Cloud Spanner.
*/
public WriteGroupedWithResults withWriteResults() {
return new WriteGroupedWithResults(spec.toBuilder().setCommitStats(true).build());
}

@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
Expand Down Expand Up @@ -1876,12 +1961,19 @@ public SpannerWriteResult expand(PCollection<MutationGroup> input) {
"Write batches to Spanner",
ParDo.of(
new WriteToSpannerFn(
spec.getSpannerConfig(), spec.getFailureMode(), FAILED_MUTATIONS_TAG))
.withOutputTags(MAIN_OUT_TAG, TupleTagList.of(FAILED_MUTATIONS_TAG)));
spec.getSpannerConfig(),
spec.getFailureMode(),
FAILED_MUTATIONS_TAG,
spec.getCommitStats(),
WRITE_RESULTS_TAG))
.withOutputTags(
MAIN_OUT_TAG, TupleTagList.of(FAILED_MUTATIONS_TAG).and(WRITE_RESULTS_TAG)));

return new SpannerWriteResult(
input.getPipeline(),
result.get(MAIN_OUT_TAG),
spec.getCommitStats() ? result.get(WRITE_RESULTS_TAG) : null,
spec.getCommitStats() ? WRITE_RESULTS_TAG : null,
result.get(FAILED_MUTATIONS_TAG),
FAILED_MUTATIONS_TAG);
}
Expand Down Expand Up @@ -2501,6 +2593,8 @@ static class GatherSortCreateBatchesFn extends DoFn<MutationGroup, Iterable<Muta
private long sortableNumCells = 0;
// total number of rows mutated in mutationsToSort
private long sortableNumRows = 0;
// minimum timestamp of elements currently in mutationsToSort
private @Nullable Instant minBundleTimestamp = null;

GatherSortCreateBatchesFn(
long maxBatchSizeBytes,
Expand Down Expand Up @@ -2529,11 +2623,13 @@ private synchronized void initSorter() {
sortableSizeBytes = 0;
sortableNumCells = 0;
sortableNumRows = 0;
minBundleTimestamp = null;
}

@FinishBundle
public synchronized void finishBundle(FinishBundleContext c) throws Exception {
sortAndOutputBatches(new OutputReceiverForFinishBundle(c));
Instant outputTimestamp = minBundleTimestamp != null ? minBundleTimestamp : Instant.now();
sortAndOutputBatches(new OutputReceiverForFinishBundle(c, outputTimestamp));
}

private synchronized void sortAndOutputBatches(OutputReceiver<Iterable<MutationGroup>> out)
Expand Down Expand Up @@ -2617,6 +2713,12 @@ public synchronized void processElement(
sortAndOutputBatches(out);
}

Instant elementTimestamp = c.timestamp();
if (elementTimestamp != null
&& (minBundleTimestamp == null || elementTimestamp.isBefore(minBundleTimestamp))) {
minBundleTimestamp = elementTimestamp;
}

mutationsToSort.add(
new MutationGroupContainer(
mg, groupSize, groupCells, groupRows, encoder.encodeTableNameAndKey(mg.primary())));
Expand Down Expand Up @@ -2663,15 +2765,15 @@ private static class OutputReceiverForFinishBundle
private final OutputBuilderSupplier outputBuilderSupplier;
private final DoFn<MutationGroup, Iterable<MutationGroup>>.FinishBundleContext context;

OutputReceiverForFinishBundle(FinishBundleContext context) {
OutputReceiverForFinishBundle(FinishBundleContext context, Instant timestamp) {
this.context = context;
this.outputBuilderSupplier =
new OutputBuilderSupplier() {
@Override
public <OutputT> WindowedValues.Builder<OutputT> builder(OutputT value) {
return WindowedValues.<OutputT>builder()
.setValue(value)
.setTimestamp(Instant.now())
.setTimestamp(timestamp)
.setPaneInfo(PaneInfo.NO_FIRING)
.setWindow(GlobalWindow.INSTANCE);
}
Expand Down Expand Up @@ -2799,16 +2901,29 @@ static class WriteToSpannerFn extends DoFn<Iterable<MutationGroup>, Void> {
Metrics.counter(WriteGrouped.class, "spanner_write_retries");

private final TupleTag<MutationGroup> failedTag;
private final boolean commitStats;
private final @Nullable TupleTag<WriteResults> writeResultsTag;

// Fluent Backoff is not serializable so create at runtime in setup().
private transient FluentBackoff bundleWriteBackoff;
private transient LoadingCache<String, ServiceCallMetric> writeMetricsByTableName;

WriteToSpannerFn(
SpannerConfig spannerConfig, FailureMode failureMode, TupleTag<MutationGroup> failedTag) {
this(spannerConfig, failureMode, failedTag, false, null);
}

WriteToSpannerFn(
SpannerConfig spannerConfig,
FailureMode failureMode,
TupleTag<MutationGroup> failedTag,
boolean commitStats,
@Nullable TupleTag<WriteResults> writeResultsTag) {
this.spannerConfig = spannerConfig;
this.failureMode = failureMode;
this.failedTag = failedTag;
this.commitStats = commitStats;
this.writeResultsTag = writeResultsTag;
}

@Setup
Expand Down Expand Up @@ -2851,9 +2966,12 @@ public void processElement(ProcessContext c) throws Exception {
mutationGroupBatchesReceived.inc();
mutationGroupsReceived.inc(mutations.size());
Iterable<Mutation> batch = Iterables.concat(mutations);
writeMutations(batch);
CommitResponse response = writeMutations(batch);
mutationGroupBatchesWriteSuccess.inc();
mutationGroupsWriteSuccess.inc(mutations.size());
if (commitStats && writeResultsTag != null) {
c.output(writeResultsTag, WriteResults.fromCommitResponse(response));
}
return;
} catch (SpannerException e) {
mutationGroupBatchesWriteFail.inc();
Expand All @@ -2872,8 +2990,11 @@ public void processElement(ProcessContext c) throws Exception {
for (MutationGroup mg : mutations) {
try {
spannerWriteRetries.inc();
writeMutations(mg);
CommitResponse response = writeMutations(mg);
mutationGroupsWriteSuccess.inc();
if (commitStats && writeResultsTag != null) {
c.output(writeResultsTag, WriteResults.fromCommitResponse(response));
}
} catch (SpannerException e) {
mutationGroupsWriteFail.inc();
LOG.warn("Failed to write the mutation group: {}", mg, e);
Expand All @@ -2886,13 +3007,15 @@ public void processElement(ProcessContext c) throws Exception {
Spanner aborts all inflight transactions during a schema change. Client is expected
to retry silently. These must not be counted against retry backoff.
*/
private void spannerWriteWithRetryIfSchemaChange(List<Mutation> batch) throws SpannerException {
private CommitResponse spannerWriteWithRetryIfSchemaChange(List<Mutation> batch)
throws SpannerException {
Set<String> tableNames = batch.stream().map(Mutation::getTable).collect(Collectors.toSet());
for (int retry = 1; ; retry++) {
try {
spannerAccessor
.getDatabaseClient()
.writeAtLeastOnceWithOptions(batch, getTransactionOptions());
CommitResponse response =
spannerAccessor
.getDatabaseClient()
.writeAtLeastOnceWithOptions(batch, getTransactionOptions());
// Get names of all tables in batch of mutations.
reportServiceCallMetricsForBatch(tableNames, "ok");
for (String tableName : tableNames) {
Expand All @@ -2906,7 +3029,7 @@ private void spannerWriteWithRetryIfSchemaChange(List<Mutation> batch) throws Sp
spannerConfig.getDatabaseId().get(),
tableName));
}
return;
return response;
} catch (AbortedException e) {
reportServiceCallMetricsForBatch(
tableNames, e.getErrorCode().getGrpcStatusCode().toString());
Expand Down Expand Up @@ -2937,7 +3060,8 @@ private Options.TransactionOption[] getTransactionOptions() {
? Options.maxCommitDelay(
java.time.Duration.ofMillis(
spannerConfig.getMaxCommitDelay().get().getMillis()))
: null)
: null,
commitStats ? Options.commitStats() : null)
.filter(Objects::nonNull)
.toArray(Options.TransactionOption[]::new);
}
Expand All @@ -2964,7 +3088,7 @@ private static ServiceCallMetric buildWriteServiceCallMetric(
}

/** Write the Mutations to Spanner, handling DEADLINE_EXCEEDED with backoff/retries. */
private void writeMutations(Iterable<Mutation> mutationIterable)
private CommitResponse writeMutations(Iterable<Mutation> mutationIterable)
throws SpannerException, IOException {
BackOff backoff = bundleWriteBackoff.backoff();
List<Mutation> mutations = ImmutableList.copyOf(mutationIterable);
Expand All @@ -2973,9 +3097,9 @@ private void writeMutations(Iterable<Mutation> mutationIterable)
Stopwatch timer = Stopwatch.createStarted();
// loop is broken on success, timeout backoff/retry attempts exceeded, or other failure.
try {
spannerWriteWithRetryIfSchemaChange(mutations);
CommitResponse response = spannerWriteWithRetryIfSchemaChange(mutations);
spannerWriteSuccess.inc();
return;
return response;
} catch (SpannerException exception) {
if (exception.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED) {
spannerWriteTimeouts.inc();
Expand Down
Loading
Loading