This is an automated email from the ASF dual-hosted git repository. pabloem pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push: new a2ebc8a95f1 Cleanup for GA launch of ReadChangeStream (#27249) a2ebc8a95f1 is described below commit a2ebc8a95f1a86d03278dd161f0d8466e4f01c4c Author: Jack Dingilian <jackdingil...@google.com> AuthorDate: Tue Jun 27 16:54:19 2023 -0400 Cleanup for GA launch of ReadChangeStream (#27249) This includes a series of small changes to release the connector: - Remove option for users to override heartbeat duration - Add option to skip creating metadata table and standalone utility to create metadata table - Add hard timeout to mutateRow requests as workaround for hanging requests - Remove excessive RCSP logging for debugging during preview - Increase read rows timeout to better accomodate large tables - Remove side effects from InitializeDoFn - Lower ReadChangeStream deadline to closer align with checkpoint duration - Add release note to CHANGES.md --- CHANGES.md | 1 + .../beam/sdk/io/gcp/bigtable/BigtableIO.java | 133 ++++++++++++++++----- .../changestreams/action/ChangeStreamAction.java | 41 +------ .../action/ReadChangeStreamPartitionAction.java | 31 +---- .../dao/BigtableChangeStreamAccessor.java | 44 +++++-- .../changestreams/dao/ChangeStreamDao.java | 14 +-- .../gcp/bigtable/changestreams/dao/DaoFactory.java | 15 ++- .../changestreams/dao/MetadataTableAdminDao.java | 5 + .../changestreams/dao/MetadataTableDao.java | 47 +++++++- .../changestreams/dofn/InitializeDoFn.java | 35 +----- .../dofn/ReadChangeStreamPartitionDoFn.java | 11 +- .../action/ChangeStreamActionTest.java | 12 +- .../ReadChangeStreamPartitionActionTest.java | 33 +++-- .../changestreams/dao/MetadataTableDaoTest.java | 31 +++++ .../changestreams/dofn/InitializeDoFnTest.java | 71 ++--------- .../dofn/ReadChangeStreamPartitionDoFnTest.java | 5 +- 16 files changed, 271 insertions(+), 258 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 6c55e276f91..aee9b96d9f4 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -59,6 +59,7 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* Support for Bigtable Change Streams added in Java `BigtableIO.ReadChangeStream` ([#27183](https://github.com/apache/beam/issues/27183)) ## New Features / Improvements diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java index 319c3139021..06497458a66 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java @@ -48,6 +48,7 @@ import org.apache.beam.sdk.io.BoundedSource.BoundedReader; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamMetrics; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.UniqueIdGenerator; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.action.ActionFactory; +import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.BigtableChangeStreamAccessor; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.DaoFactory; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.MetadataTableAdminDao; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dofn.DetectNewPartitionsDoFn; @@ -310,7 +311,6 @@ public class BigtableIO { * * <ul> * <li>{@link BigtableIO.ReadChangeStream#withStartTime} which defaults to now. - * <li>{@link BigtableIO.ReadChangeStream#withHeartbeatDuration} with defaults to 1 seconds. * <li>{@link BigtableIO.ReadChangeStream#withMetadataTableProjectId} which defaults to value * from {@link BigtableIO.ReadChangeStream#withProjectId} * <li>{@link BigtableIO.ReadChangeStream#withMetadataTableInstanceId} which defaults to value @@ -1797,8 +1797,6 @@ public class BigtableIO { RESUME_OR_NEW, // Same as RESUME_OR_NEW except if previous pipeline doesn't exist, don't start. RESUME_OR_FAIL, - // Start a new pipeline. Overriding existing pipeline with the same name. - NEW, // This skips cleaning up previous pipeline metadata and starts a new pipeline. This should // only be used to skip cleanup in tests @VisibleForTesting @@ -1827,8 +1825,6 @@ public class BigtableIO { abstract @Nullable Instant getEndTime(); - abstract @Nullable Duration getHeartbeatDuration(); - abstract @Nullable String getChangeStreamName(); abstract @Nullable ExistingPipelineOptions getExistingPipelineOptions(); @@ -1837,6 +1833,8 @@ public class BigtableIO { abstract @Nullable String getMetadataTableId(); + abstract @Nullable Boolean getCreateOrUpdateMetadataTable(); + abstract ReadChangeStream.Builder toBuilder(); /** @@ -1909,16 +1907,6 @@ public class BigtableIO { return toBuilder().setEndTime(endTime).build(); } - /** - * Returns a new {@link BigtableIO.ReadChangeStream} that will send heartbeat messages at - * specified interval. - * - * <p>Does not modify this object. - */ - public ReadChangeStream withHeartbeatDuration(Duration interval) { - return toBuilder().setHeartbeatDuration(interval).build(); - } - /** * Returns a new {@link BigtableIO.ReadChangeStream} that uses changeStreamName as prefix for * the metadata table. @@ -2000,6 +1988,19 @@ public class BigtableIO { .build(); } + /** + * Returns a new {@link BigtableIO.ReadChangeStream} that, if set to true, will create or update + * metadata table before launching pipeline. Otherwise, it is expected that a metadata table + * with correct schema exists. + * + * <p>Optional: defaults to true + * + * <p>Does not modify this object. + */ + public ReadChangeStream withCreateOrUpdateMetadataTable(boolean shouldCreate) { + return toBuilder().setCreateOrUpdateMetadataTable(shouldCreate).build(); + } + @Override public PCollection<KV<ByteString, ChangeStreamMutation>> expand(PBegin input) { checkArgument( @@ -2040,10 +2041,6 @@ public class BigtableIO { if (startTime == null) { startTime = Instant.now(); } - Duration heartbeatDuration = getHeartbeatDuration(); - if (heartbeatDuration == null) { - heartbeatDuration = Duration.standardSeconds(1); - } String changeStreamName = getChangeStreamName(); if (changeStreamName == null || changeStreamName.isEmpty()) { changeStreamName = UniqueIdGenerator.generateRowKeyPrefix(); @@ -2053,21 +2050,55 @@ public class BigtableIO { existingPipelineOptions = ExistingPipelineOptions.FAIL_IF_EXISTS; } + boolean shouldCreateOrUpdateMetadataTable = true; + if (getCreateOrUpdateMetadataTable() != null) { + shouldCreateOrUpdateMetadataTable = getCreateOrUpdateMetadataTable(); + } + ActionFactory actionFactory = new ActionFactory(); + ChangeStreamMetrics metrics = new ChangeStreamMetrics(); DaoFactory daoFactory = new DaoFactory( bigtableConfig, metadataTableConfig, getTableId(), metadataTableId, changeStreamName); - ChangeStreamMetrics metrics = new ChangeStreamMetrics(); + + try { + MetadataTableAdminDao metadataTableAdminDao = daoFactory.getMetadataTableAdminDao(); + checkArgument(metadataTableAdminDao != null); + checkArgument( + metadataTableAdminDao.isAppProfileSingleClusterAndTransactional( + metadataTableConfig.getAppProfileId().get()), + "App profile id '" + + metadataTableConfig.getAppProfileId().get() + + "' provided to access metadata table needs to use single-cluster routing policy" + + " and allow single-row transactions."); + + // Only try to create or update metadata table if option is set to true. Otherwise, just + // check if the table exists. + if (shouldCreateOrUpdateMetadataTable && metadataTableAdminDao.createMetadataTable()) { + LOG.info("Created metadata table: " + metadataTableAdminDao.getTableId()); + } + checkArgument( + metadataTableAdminDao.doesMetadataTableExist(), + "Metadata table does not exist: " + metadataTableAdminDao.getTableId()); + + try (BigtableChangeStreamAccessor bigtableChangeStreamAccessor = + BigtableChangeStreamAccessor.getOrCreate(bigtableConfig)) { + checkArgument( + bigtableChangeStreamAccessor.getTableAdminClient().exists(getTableId()), + "Change Stream table does not exist"); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + daoFactory.close(); + } + InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, - metadataTableConfig.getAppProfileId().get(), - startTime, - existingPipelineOptions); + new InitializeDoFn(daoFactory, startTime, existingPipelineOptions); DetectNewPartitionsDoFn detectNewPartitionsDoFn = new DetectNewPartitionsDoFn(getEndTime(), actionFactory, daoFactory, metrics); ReadChangeStreamPartitionDoFn readChangeStreamPartitionDoFn = - new ReadChangeStreamPartitionDoFn(heartbeatDuration, daoFactory, actionFactory, metrics); + new ReadChangeStreamPartitionDoFn(daoFactory, actionFactory, metrics); PCollection<KV<ByteString, ChangeStreamRecord>> readChangeStreamOutput = input @@ -2101,14 +2132,60 @@ public class BigtableIO { abstract ReadChangeStream.Builder setEndTime(Instant endTime); - abstract ReadChangeStream.Builder setHeartbeatDuration(Duration interval); - abstract ReadChangeStream.Builder setChangeStreamName(String changeStreamName); abstract ReadChangeStream.Builder setExistingPipelineOptions( ExistingPipelineOptions existingPipelineOptions); + abstract ReadChangeStream.Builder setCreateOrUpdateMetadataTable(boolean shouldCreate); + abstract ReadChangeStream build(); } } + + /** + * Utility method to create or update Read Change Stream metadata table. This requires Bigtable + * table create permissions. This method is useful if the pipeline isn't granted permissions to + * create Bigtable tables. Run this method with correct permissions to create the metadata table, + * which is required to read Bigtable change streams. This method only needs to be run once, and + * the metadata table can be reused for all pipelines. + * + * @param projectId project id of the metadata table, usually the same as the project of the table + * being streamed + * @param instanceId instance id of the metadata table, usually the same as the instance of the + * table being streamed + * @param tableId name of the metadata table, leave it null or empty to use default. + * @return true if the table was successfully created. Otherwise, false. + */ + public static boolean createOrUpdateReadChangeStreamMetadataTable( + String projectId, String instanceId, @Nullable String tableId) throws IOException { + BigtableConfig bigtableConfig = + BigtableConfig.builder() + .setValidate(true) + .setProjectId(StaticValueProvider.of(projectId)) + .setInstanceId(StaticValueProvider.of(instanceId)) + .setAppProfileId( + StaticValueProvider.of( + "default")) // App profile is not used. It's only required for data API. + .build(); + + if (tableId == null || tableId.isEmpty()) { + tableId = MetadataTableAdminDao.DEFAULT_METADATA_TABLE_NAME; + } + + DaoFactory daoFactory = new DaoFactory(null, bigtableConfig, null, tableId, null); + + try { + MetadataTableAdminDao metadataTableAdminDao = daoFactory.getMetadataTableAdminDao(); + + // Only try to create or update metadata table if option is set to true. Otherwise, just + // check if the table exists. + if (metadataTableAdminDao.createMetadataTable()) { + LOG.info("Created metadata table: " + metadataTableAdminDao.getTableId()); + } + return metadataTableAdminDao.doesMetadataTableExist(); + } finally { + daoFactory.close(); + } + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java index 3ab3b1973b9..adb817a4c08 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java @@ -28,7 +28,6 @@ import com.google.cloud.bigtable.data.v2.models.Heartbeat; import com.google.cloud.bigtable.data.v2.models.Range; import com.google.protobuf.ByteString; import java.util.Optional; -import java.util.stream.Collectors; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamMetrics; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.estimator.BytesThroughputEstimator; @@ -108,8 +107,7 @@ public class ChangeStreamAction { RestrictionTracker<StreamProgress, StreamProgress> tracker, DoFn.OutputReceiver<KV<ByteString, ChangeStreamRecord>> receiver, ManualWatermarkEstimator<Instant> watermarkEstimator, - BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator, - boolean shouldDebug) { + BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator) { if (record instanceof Heartbeat) { Heartbeat heartbeat = (Heartbeat) record; final Instant watermark = toJodaTime(heartbeat.getEstimatedLowWatermark()); @@ -129,24 +127,11 @@ public class ChangeStreamAction { true); watermarkEstimator.setWatermark(watermark); - if (shouldDebug) { - LOG.info( - "RCSP {}: Heartbeat partition: {} token: {} watermark: {}", - formatByteStringRange(partitionRecord.getPartition()), - formatByteStringRange(heartbeat.getChangeStreamContinuationToken().getPartition()), - heartbeat.getChangeStreamContinuationToken().getToken(), - heartbeat.getEstimatedLowWatermark()); - } // If the tracker fail to claim the streamProgress, it most likely means the runner initiated // a checkpoint. See {@link // org.apache.beam.sdk.io.gcp.bigtable.changestreams.restriction.ReadChangeStreamPartitionProgressTracker} // for more information regarding runner initiated checkpoints. if (!tracker.tryClaim(streamProgress)) { - if (shouldDebug) { - LOG.info( - "RCSP {}: Checkpoint heart beat tracker", - formatByteStringRange(partitionRecord.getPartition())); - } return Optional.of(DoFn.ProcessContinuation.stop()); } metrics.incHeartbeatCount(); @@ -163,30 +148,11 @@ public class ChangeStreamAction { CloseStream closeStream = (CloseStream) record; StreamProgress streamProgress = new StreamProgress(closeStream); - if (shouldDebug) { - LOG.info( - "RCSP {}: CloseStream: {}", - formatByteStringRange(partitionRecord.getPartition()), - closeStream.getChangeStreamContinuationTokens().stream() - .map( - c -> - "{partition: " - + formatByteStringRange(c.getPartition()) - + " token: " - + c.getToken() - + "}") - .collect(Collectors.joining(", ", "[", "]"))); - } // If the tracker fail to claim the streamProgress, it most likely means the runner initiated // a checkpoint. See {@link // org.apache.beam.sdk.io.gcp.bigtable.changestreams.restriction.ReadChangeStreamPartitionProgressTracker} // for more information regarding runner initiated checkpoints. if (!tracker.tryClaim(streamProgress)) { - if (shouldDebug) { - LOG.info( - "RCSP {}: Checkpoint close stream tracker", - formatByteStringRange(partitionRecord.getPartition())); - } return Optional.of(DoFn.ProcessContinuation.stop()); } metrics.incClosestreamCount(); @@ -217,11 +183,6 @@ public class ChangeStreamAction { // a checkpoint. See ReadChangeStreamPartitionProgressTracker for more information regarding // runner initiated checkpoints. if (!tracker.tryClaim(streamProgress)) { - if (shouldDebug) { - LOG.info( - "RCSP {}: Checkpoint data change tracker", - formatByteStringRange(partitionRecord.getPartition())); - } return Optional.of(DoFn.ProcessContinuation.stop()); } if (changeStreamMutation.getType() == ChangeStreamMutation.MutationType.GARBAGE_COLLECTION) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java index 87273c9894d..8638339ac5c 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java @@ -138,23 +138,8 @@ public class ReadChangeStreamPartitionAction { OutputReceiver<KV<ByteString, ChangeStreamRecord>> receiver, ManualWatermarkEstimator<Instant> watermarkEstimator) throws IOException { - // Watermark being delayed beyond 5 minutes signals a possible problem. - boolean shouldDebug = - watermarkEstimator.getState().plus(Duration.standardMinutes(5)).isBeforeNow(); BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator = new BytesThroughputEstimator<>(sizeEstimator, Instant.now()); - - if (shouldDebug) { - LOG.info( - "RCSP {}: Partition: " - + partitionRecord - + "\n Watermark: " - + watermarkEstimator.getState() - + "\n RestrictionTracker: " - + tracker.currentRestriction(), - formatByteStringRange(partitionRecord.getPartition())); - } - // Lock the partition if (tracker.currentRestriction().isEmpty()) { boolean lockedPartition = metadataTableDao.lockAndRecordPartition(partitionRecord); @@ -266,12 +251,10 @@ public class ReadChangeStreamPartitionAction { new NewPartition( childPartition, Collections.singletonList(token), watermarkEstimator.getState())); } - if (shouldDebug) { - LOG.info( - "RCSP {}: Split/Merge into {}", - formatByteStringRange(partitionRecord.getPartition()), - partitionsToString(childPartitions)); - } + LOG.info( + "RCSP {}: Split/Merge into {}", + formatByteStringRange(partitionRecord.getPartition()), + partitionsToString(childPartitions)); if (!coverSameKeySpace(tokenPartitions, partitionRecord.getPartition())) { LOG.warn( "RCSP {}: CloseStream has tokens {} that don't cover the entire keyspace", @@ -299,8 +282,7 @@ public class ReadChangeStreamPartitionAction { partitionRecord, tracker.currentRestriction(), partitionRecord.getEndTime(), - heartbeatDuration, - shouldDebug); + heartbeatDuration); for (ChangeStreamRecord record : stream) { Optional<ProcessContinuation> result = changeStreamAction.run( @@ -309,8 +291,7 @@ public class ReadChangeStreamPartitionAction { tracker, receiver, watermarkEstimator, - throughputEstimator, - shouldDebug); + throughputEstimator); // changeStreamAction will usually return Optional.empty() except for when a checkpoint // (either runner or pipeline initiated) is required. if (result.isPresent()) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java index 69c6efba1de..ecf9a703959 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java @@ -47,7 +47,8 @@ import org.threeten.bp.Duration; * backend and the jobs on the same machine shares the same sets of connections. */ @Internal -class BigtableChangeStreamAccessor { +public class BigtableChangeStreamAccessor implements AutoCloseable { + static final Duration MUTATE_ROW_DEADLINE = Duration.ofSeconds(30); // Create one bigtable data/admin client per bigtable config (project/instance/table/app profile) private static final ConcurrentHashMap<BigtableConfig, BigtableChangeStreamAccessor> bigtableAccessors = new ConcurrentHashMap<>(); @@ -55,14 +56,31 @@ class BigtableChangeStreamAccessor { private final BigtableDataClient dataClient; private final BigtableTableAdminClient tableAdminClient; private final BigtableInstanceAdminClient instanceAdminClient; + private final BigtableConfig bigtableConfig; private BigtableChangeStreamAccessor( BigtableDataClient dataClient, BigtableTableAdminClient tableAdminClient, - BigtableInstanceAdminClient instanceAdminClient) { + BigtableInstanceAdminClient instanceAdminClient, + BigtableConfig bigtableConfig) { this.dataClient = dataClient; this.tableAdminClient = tableAdminClient; this.instanceAdminClient = instanceAdminClient; + this.bigtableConfig = bigtableConfig; + } + + @Override + public synchronized void close() { + if (dataClient != null) { + dataClient.close(); + } + if (tableAdminClient != null) { + tableAdminClient.close(); + } + if (instanceAdminClient != null) { + instanceAdminClient.close(); + } + bigtableAccessors.remove(bigtableConfig); } /** @@ -146,9 +164,10 @@ class BigtableChangeStreamAccessor { .readRowsSettings() .setRetrySettings( readRowsRetrySettings - .setInitialRpcTimeout(Duration.ofSeconds(30)) - .setTotalTimeout(Duration.ofSeconds(30)) - .setMaxRpcTimeout(Duration.ofSeconds(30)) + // metadata table scans can get quite large, so use a higher deadline + .setInitialRpcTimeout(Duration.ofMinutes(3)) + .setTotalTimeout(Duration.ofMinutes(3)) + .setMaxRpcTimeout(Duration.ofMinutes(3)) .setMaxAttempts(10) .build()); @@ -159,9 +178,9 @@ class BigtableChangeStreamAccessor { .mutateRowSettings() .setRetrySettings( mutateRowRetrySettings - .setInitialRpcTimeout(Duration.ofSeconds(30)) - .setTotalTimeout(Duration.ofSeconds(30)) - .setMaxRpcTimeout(Duration.ofSeconds(30)) + .setInitialRpcTimeout(MUTATE_ROW_DEADLINE) + .setTotalTimeout(MUTATE_ROW_DEADLINE) + .setMaxRpcTimeout(MUTATE_ROW_DEADLINE) .setMaxAttempts(10) .build()); @@ -185,9 +204,9 @@ class BigtableChangeStreamAccessor { .readChangeStreamSettings() .setRetrySettings( readChangeStreamRetrySettings - .setInitialRpcTimeout(Duration.ofSeconds(60)) - .setTotalTimeout(Duration.ofSeconds(60)) - .setMaxRpcTimeout(Duration.ofSeconds(60)) + .setInitialRpcTimeout(Duration.ofSeconds(15)) + .setTotalTimeout(Duration.ofSeconds(15)) + .setMaxRpcTimeout(Duration.ofSeconds(15)) .setMaxAttempts(10) .build()); @@ -196,7 +215,8 @@ class BigtableChangeStreamAccessor { BigtableTableAdminClient.create(tableAdminSettingsBuilder.build()); BigtableInstanceAdminClient instanceAdminClient = BigtableInstanceAdminClient.create(instanceAdminSettingsBuilder.build()); - return new BigtableChangeStreamAccessor(dataClient, tableAdminClient, instanceAdminClient); + return new BigtableChangeStreamAccessor( + dataClient, tableAdminClient, instanceAdminClient, bigtableConfig); } public BigtableDataClient getDataClient() { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java index 1c24e09db7d..5974af22f86 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java @@ -17,7 +17,6 @@ */ package org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao; -import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.ByteStringRangeHelper.formatByteStringRange; import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.TimestampConverter.toThreetenInstant; import com.google.api.gax.rpc.ServerStream; @@ -36,14 +35,10 @@ import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.PartitionRecord; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.restriction.StreamProgress; import org.joda.time.Duration; import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; /** Data access object to list and read stream partitions of a table. */ @Internal public class ChangeStreamDao { - private static final Logger LOG = LoggerFactory.getLogger(ChangeStreamDao.class); - private final BigtableDataClient dataClient; private final String tableId; @@ -75,8 +70,7 @@ public class ChangeStreamDao { PartitionRecord partition, StreamProgress streamProgress, @Nullable Instant endTime, - Duration heartbeatDuration, - boolean shouldDebug) + Duration heartbeatDuration) throws IOException { ReadChangeStreamQuery query = ReadChangeStreamQuery.create(tableId).streamPartition(partition.getPartition()); @@ -98,12 +92,6 @@ public class ChangeStreamDao { query.endTime(TimestampConverter.toThreetenInstant(endTime)); } query.heartbeatDuration(org.threeten.bp.Duration.ofMillis(heartbeatDuration.getMillis())); - if (shouldDebug) { - LOG.info( - "RCSP {} ReadChangeStreamRequest: {}", - formatByteStringRange(partition.getPartition()), - query); - } return dataClient.readChangeStream(query); } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java index acecc22b891..35ac8ed29c3 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java @@ -30,7 +30,7 @@ import org.apache.beam.sdk.io.gcp.bigtable.BigtableConfig; // Allows transient fields to be intialized later @SuppressWarnings("initialization.fields.uninitialized") @Internal -public class DaoFactory implements Serializable { +public class DaoFactory implements Serializable, AutoCloseable { private static final long serialVersionUID = 3732208768248394205L; private transient ChangeStreamDao changeStreamDao; @@ -58,6 +58,19 @@ public class DaoFactory implements Serializable { this.metadataTableId = metadataTableId; } + @Override + public void close() { + try { + if (metadataTableAdminDao != null || metadataTableDao != null) { + BigtableChangeStreamAccessor.getOrCreate(metadataTableConfig).close(); + } + if (changeStreamDao != null) { + BigtableChangeStreamAccessor.getOrCreate(changeStreamConfig).close(); + } + } catch (Exception ignored) { + } + } + public String getChangeStreamName() { return changeStreamName; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java index 5194ca49d62..73bfbe5865b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java @@ -128,6 +128,11 @@ public class MetadataTableAdminDao { return false; } + /** @return true if metadata table exists, otherwise false. */ + public boolean doesMetadataTableExist() { + return tableAdminClient.exists(tableId); + } + /** * Create the metadata table if it does not exist yet. If the table does exist, verify all the * column families exists, if not add those column families. This table only need to be created diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java index d4938e6bef2..418c29d1de7 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java @@ -30,6 +30,7 @@ import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.Metadata import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder.parseWatermarkFromRow; import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder.parseWatermarkLastUpdatedFromRow; +import com.google.api.core.ApiFuture; import com.google.api.gax.rpc.ServerStream; import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.models.ChangeStreamContinuationToken; @@ -50,6 +51,9 @@ import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import javax.annotation.Nonnull; import javax.annotation.Nullable; import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationException; @@ -60,6 +64,7 @@ import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.DetectNewPartitio import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.NewPartition; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.PartitionRecord; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.StreamPartitionWithWatermark; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.Longs; import org.joda.time.Duration; import org.joda.time.Instant; @@ -319,7 +324,7 @@ public class MetadataTableDao { .deleteCells( MetadataTableAdminDao.CF_SHOULD_DELETE, ByteStringRange.serializeToByteString(parentPartition)); - dataClient.mutateRow(rowMutation); + mutateRowWithHardTimeout(rowMutation); } /** @@ -352,7 +357,7 @@ public class MetadataTableDao { ByteStringRange.serializeToByteString(token.getPartition()), 1); } - dataClient.mutateRow(rowMutation); + mutateRowWithHardTimeout(rowMutation); } /** @@ -511,7 +516,7 @@ public class MetadataTableDao { MetadataTableAdminDao.QUALIFIER_DEFAULT, currentToken.getToken()); } - dataClient.mutateRow(rowMutation); + mutateRowWithHardTimeout(rowMutation); } /** @@ -705,6 +710,10 @@ public class MetadataTableDao { boolean lockAcquired = !dataClient.checkAndMutateRow(rowMutation); if (lockAcquired) { + LOG.info( + "RCSP: {} acquired lock for uid: {}", + formatByteStringRange(partitionRecord.getPartition()), + partitionRecord.getUuid()); return true; } else { // If the lock is already held we need to check if it was acquired by a duplicate @@ -724,7 +733,7 @@ public class MetadataTableDao { MetadataTableAdminDao.CF_VERSION, MetadataTableAdminDao.QUALIFIER_DEFAULT, MetadataTableAdminDao.CURRENT_METADATA_TABLE_VERSION); - dataClient.mutateRow(rowMutation); + mutateRowWithHardTimeout(rowMutation); } /** @@ -778,6 +787,34 @@ public class MetadataTableDao { MetadataTableAdminDao.CF_MISSING_PARTITIONS, ByteString.copyFromUtf8(MetadataTableAdminDao.QUALIFIER_DEFAULT), ByteString.copyFrom(serializedMissingPartition)); - dataClient.mutateRow(rowMutation); + mutateRowWithHardTimeout(rowMutation); + } + + /** + * This adds a hard timeout of 40 seconds to mutate row futures. These requests already have a + * 30-second deadline. This is a workaround for an extremely rare issue we see with requests not + * respecting their deadlines. This can be removed once we've pinpointed the cause. + * + * @param rowMutation Bigtable RowMutation to apply + */ + @VisibleForTesting + void mutateRowWithHardTimeout(RowMutation rowMutation) { + ApiFuture<Void> mutateRowFuture = dataClient.mutateRowAsync(rowMutation); + try { + mutateRowFuture.get( + BigtableChangeStreamAccessor.MUTATE_ROW_DEADLINE.getSeconds() + 10, TimeUnit.SECONDS); + } catch (TimeoutException timeoutException) { + mutateRowFuture.cancel(true); + throw new RuntimeException( + "Cancelled mutateRow request after exceeding deadline", timeoutException); + } catch (ExecutionException executionException) { + if (executionException.getCause() instanceof RuntimeException) { + throw (RuntimeException) executionException.getCause(); + } + throw new RuntimeException(executionException); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw new RuntimeException(interruptedException); + } } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java index 07826ee1db7..c8ee12772cc 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java @@ -36,20 +36,14 @@ import org.slf4j.LoggerFactory; @Internal public class InitializeDoFn extends DoFn<byte[], InitialPipelineState> implements Serializable { private static final long serialVersionUID = 1868189906451252363L; - private static final Logger LOG = LoggerFactory.getLogger(InitializeDoFn.class); private final DaoFactory daoFactory; - private final String metadataTableAppProfileId; private Instant startTime; private final ExistingPipelineOptions existingPipelineOptions; public InitializeDoFn( - DaoFactory daoFactory, - String metadataTableAppProfileId, - Instant startTime, - ExistingPipelineOptions existingPipelineOptions) { + DaoFactory daoFactory, Instant startTime, ExistingPipelineOptions existingPipelineOptions) { this.daoFactory = daoFactory; - this.metadataTableAppProfileId = metadataTableAppProfileId; this.startTime = startTime; this.existingPipelineOptions = existingPipelineOptions; } @@ -59,35 +53,12 @@ public class InitializeDoFn extends DoFn<byte[], InitialPipelineState> implement LOG.info(daoFactory.getStreamTableDebugString()); LOG.info(daoFactory.getMetadataTableDebugString()); LOG.info("ChangeStreamName: " + daoFactory.getChangeStreamName()); - if (!daoFactory - .getMetadataTableAdminDao() - .isAppProfileSingleClusterAndTransactional(this.metadataTableAppProfileId)) { - LOG.error( - "App profile id '" - + metadataTableAppProfileId - + "' provided to access metadata table needs to use single-cluster routing policy" - + " and allow single-row transactions."); - // Terminate this pipeline now. - return; - } - if (daoFactory.getMetadataTableAdminDao().createMetadataTable()) { - LOG.info("Created metadata table: " + daoFactory.getMetadataTableAdminDao().getTableId()); - } else { - LOG.info( - "Reusing existing metadata table: " + daoFactory.getMetadataTableAdminDao().getTableId()); - } boolean resume = false; DetectNewPartitionsState detectNewPartitionsState = daoFactory.getMetadataTableDao().readDetectNewPartitionsState(); switch (existingPipelineOptions) { - case NEW: - // clean up table - LOG.info( - "Cleaning up an old pipeline with the same change stream name to start a new pipeline with the same name."); - daoFactory.getMetadataTableAdminDao().cleanUpPrefix(); - break; case RESUME_OR_NEW: // perform resumption. if (detectNewPartitionsState != null) { @@ -98,7 +69,6 @@ public class InitializeDoFn extends DoFn<byte[], InitialPipelineState> implement LOG.info( "Attempted to resume, but previous watermark does not exist, starting at {}", startTime); - daoFactory.getMetadataTableAdminDao().cleanUpPrefix(); } break; case RESUME_OR_FAIL: @@ -118,9 +88,6 @@ public class InitializeDoFn extends DoFn<byte[], InitialPipelineState> implement "A previous pipeline exists with the same change stream name and existingPipelineOption is set to FAIL_IF_EXISTS."); return; } - // We still want to clean up any existing prefixes in case there are lingering metadata that - // would interfere with the new run. - daoFactory.getMetadataTableAdminDao().cleanUpPrefix(); break; case SKIP_CLEANUP: if (detectNewPartitionsState != null) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java index 91210caa597..e4c911f926a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java @@ -58,9 +58,8 @@ public class ReadChangeStreamPartitionDoFn private static final long serialVersionUID = 4418739381635104479L; private static final BigDecimal MAX_DOUBLE = BigDecimal.valueOf(Double.MAX_VALUE); private static final Logger LOG = LoggerFactory.getLogger(ReadChangeStreamPartitionDoFn.class); - public static final int THROUGHPUT_ESTIMATION_WINDOW_SECONDS = 10; + private static final Duration HEARTBEAT_DURATION = Duration.standardSeconds(1); - private final Duration heartbeatDuration; private final DaoFactory daoFactory; private final ChangeStreamMetrics metrics; private final ActionFactory actionFactory; @@ -68,11 +67,7 @@ public class ReadChangeStreamPartitionDoFn private ReadChangeStreamPartitionAction readChangeStreamPartitionAction; public ReadChangeStreamPartitionDoFn( - Duration heartbeatDuration, - DaoFactory daoFactory, - ActionFactory actionFactory, - ChangeStreamMetrics metrics) { - this.heartbeatDuration = heartbeatDuration; + DaoFactory daoFactory, ActionFactory actionFactory, ChangeStreamMetrics metrics) { this.daoFactory = daoFactory; this.metrics = metrics; this.actionFactory = actionFactory; @@ -171,7 +166,7 @@ public class ReadChangeStreamPartitionDoFn changeStreamDao, metrics, changeStreamAction, - heartbeatDuration, + HEARTBEAT_DURATION, sizeEstimator); } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java index 2e340c2b30b..b2c9d6ac131 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java @@ -103,8 +103,7 @@ public class ChangeStreamActionTest { tracker, receiver, watermarkEstimator, - throughputEstimator, - false); + throughputEstimator); assertFalse(result.isPresent()); verify(metrics).incHeartbeatCount(); @@ -152,8 +151,7 @@ public class ChangeStreamActionTest { tracker, receiver, watermarkEstimator, - throughputEstimator, - false); + throughputEstimator); assertTrue(result.isPresent()); assertEquals(DoFn.ProcessContinuation.resume(), result.get()); @@ -187,8 +185,7 @@ public class ChangeStreamActionTest { tracker, receiver, watermarkEstimator, - throughputEstimator, - false); + throughputEstimator); assertFalse(result.isPresent()); verify(metrics).incChangeStreamMutationUserCounter(); @@ -243,8 +240,7 @@ public class ChangeStreamActionTest { tracker, receiver, watermarkEstimator, - throughputEstimator, - false); + throughputEstimator); assertFalse(result.isPresent()); verify(metrics).incChangeStreamMutationGcCounter(); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java index 305be156885..43419b7147f 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java @@ -20,7 +20,6 @@ package org.apache.beam.sdk.io.gcp.bigtable.changestreams.action; import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamContinuationTokenHelper.getTokenWithCorrectPartition; import static org.junit.Assert.assertEquals; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -128,10 +127,10 @@ public class ReadChangeStreamPartitionActionTest { Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class); when(responseIterator.next()).thenReturn(mockHeartBeat); when(responseIterator.hasNext()).thenReturn(true); - when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean())) + when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any())) .thenReturn(responses); - when(changeStreamAction.run(any(), any(), any(), any(), any(), any(), anyBoolean())) + when(changeStreamAction.run(any(), any(), any(), any(), any(), any())) .thenReturn(Optional.of(DoFn.ProcessContinuation.stop())); final DoFn.ProcessContinuation result = @@ -140,7 +139,7 @@ public class ReadChangeStreamPartitionActionTest { assertEquals(DoFn.ProcessContinuation.stop(), result); // Verify that on successful lock, we don't tryClaim on the tracker. verify(tracker, never()).tryClaim(any()); - verify(changeStreamAction).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction).run(any(), any(), any(), any(), any(), any()); } @Test @@ -157,7 +156,7 @@ public class ReadChangeStreamPartitionActionTest { StreamProgress streamProgress = new StreamProgress(); streamProgress.setFailToLock(true); verify(tracker).tryClaim(streamProgress); - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); } @Test @@ -174,10 +173,10 @@ public class ReadChangeStreamPartitionActionTest { Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class); when(responseIterator.next()).thenReturn(mockHeartBeat); when(responseIterator.hasNext()).thenReturn(true); - when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean())) + when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any())) .thenReturn(responses); - when(changeStreamAction.run(any(), any(), any(), any(), any(), any(), anyBoolean())) + when(changeStreamAction.run(any(), any(), any(), any(), any(), any())) .thenReturn(Optional.of(DoFn.ProcessContinuation.stop())); final DoFn.ProcessContinuation result = @@ -186,7 +185,7 @@ public class ReadChangeStreamPartitionActionTest { assertEquals(DoFn.ProcessContinuation.stop(), result); // Verify that on successful lock, we don't tryClaim on the tracker. verify(tracker, never()).tryClaim(any()); - verify(changeStreamAction).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction).run(any(), any(), any(), any(), any(), any()); } @Test @@ -205,7 +204,7 @@ public class ReadChangeStreamPartitionActionTest { StreamProgress streamProgress = new StreamProgress(); streamProgress.setFailToLock(true); verify(tracker).tryClaim(streamProgress); - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); } @Test @@ -221,16 +220,16 @@ public class ReadChangeStreamPartitionActionTest { Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class); when(responseIterator.next()).thenReturn(mockHeartBeat); when(responseIterator.hasNext()).thenReturn(true); - when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean())) + when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any())) .thenReturn(responses); - when(changeStreamAction.run(any(), any(), any(), any(), any(), any(), anyBoolean())) + when(changeStreamAction.run(any(), any(), any(), any(), any(), any())) .thenReturn(Optional.of(DoFn.ProcessContinuation.stop())); final DoFn.ProcessContinuation result = action.run(partitionRecord, tracker, receiver, watermarkEstimator); assertEquals(DoFn.ProcessContinuation.stop(), result); - verify(changeStreamAction).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction).run(any(), any(), any(), any(), any(), any()); } @Test @@ -247,7 +246,7 @@ public class ReadChangeStreamPartitionActionTest { action.run(partitionRecord, tracker, receiver, watermarkEstimator); assertEquals(DoFn.ProcessContinuation.stop(), result); // Should terminate before reaching processing stream partition responses. - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); // Should not try claim any restriction when processing CloseStream verify(tracker, (never())).tryClaim(any()); // Should decrement the metric on termination. @@ -273,7 +272,7 @@ public class ReadChangeStreamPartitionActionTest { action.run(partitionRecord, tracker, receiver, watermarkEstimator); assertEquals(DoFn.ProcessContinuation.stop(), result); // Should terminate before reaching processing stream partition responses. - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); // Should not try claim any restriction when processing CloseStream verify(tracker, (never())).tryClaim(any()); // Should decrement the metric on termination. @@ -308,7 +307,7 @@ public class ReadChangeStreamPartitionActionTest { action.run(partitionRecord, tracker, receiver, watermarkEstimator); assertEquals(DoFn.ProcessContinuation.stop(), result); // Should terminate before reaching processing stream partition responses. - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); // Should not try claim any restriction when processing CloseStream verify(tracker, (never())).tryClaim(any()); // Should decrement the metric on termination. @@ -360,7 +359,7 @@ public class ReadChangeStreamPartitionActionTest { action.run(partitionRecord, tracker, receiver, watermarkEstimator); assertEquals(DoFn.ProcessContinuation.stop(), result); // Should terminate before reaching processing stream partition responses. - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); // Should not try claim any restriction when processing CloseStream verify(tracker, (never())).tryClaim(any()); // Should decrement the metric on termination. @@ -406,7 +405,7 @@ public class ReadChangeStreamPartitionActionTest { action.run(partitionRecord, tracker, receiver, watermarkEstimator); assertEquals(DoFn.ProcessContinuation.stop(), result); // Should terminate before reaching processing stream partition responses. - verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean()); + verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any()); // Should not try claim any restriction when processing CloseStream verify(tracker, (never())).tryClaim(any()); // Should decrement the metric on termination. diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java index c82e3eac299..867117b4d39 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java @@ -23,10 +23,14 @@ import static org.hamcrest.Matchers.containsInAnyOrder; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; 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.when; +import com.google.api.core.ApiFuture; import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient; import com.google.cloud.bigtable.admin.v2.BigtableTableAdminSettings; import com.google.cloud.bigtable.data.v2.BigtableDataClient; @@ -45,6 +49,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.UniqueIdGenerator; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder; import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.NewPartition; @@ -743,4 +750,28 @@ public class MetadataTableDaoTest { .get(0) .getValue()); } + + @Test + public void mutateRowWithHardTimeoutErrorHandling() + throws ExecutionException, InterruptedException, TimeoutException { + BigtableDataClient mockClient = Mockito.mock(BigtableDataClient.class); + MetadataTableDao daoWithMock = + new MetadataTableDao(mockClient, "test-table", ByteString.copyFromUtf8("test")); + ApiFuture<Void> mockFuture = mock(ApiFuture.class); + when(mockClient.mutateRowAsync(any())).thenReturn(mockFuture); + + when(mockFuture.get(40, TimeUnit.SECONDS)) + .thenThrow(TimeoutException.class) + .thenThrow(InterruptedException.class) + .thenThrow(ExecutionException.class); + assertThrows( + RuntimeException.class, + () -> daoWithMock.mutateRowWithHardTimeout(RowMutation.create("test", "test").deleteRow())); + assertThrows( + RuntimeException.class, + () -> daoWithMock.mutateRowWithHardTimeout(RowMutation.create("test", "test").deleteRow())); + assertThrows( + RuntimeException.class, + () -> daoWithMock.mutateRowWithHardTimeout(RowMutation.create("test", "test").deleteRow())); + } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java index a9fa54d4c49..f0337504282 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java @@ -20,7 +20,6 @@ package org.apache.beam.sdk.io.gcp.bigtable.changestreams.dofn; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.never; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; @@ -88,8 +87,6 @@ public class InitializeDoFnTest { metadataTableAdminDao = spy(new MetadataTableAdminDao(adminClient, null, changeStreamName, tableId)); metadataTableAdminDao.createMetadataTable(); - doReturn(true).when(metadataTableAdminDao).isAppProfileSingleClusterAndTransactional(any()); - when(daoFactory.getMetadataTableAdminDao()).thenReturn(metadataTableAdminDao); metadataTableDao = new MetadataTableDao( dataClient, tableId, metadataTableAdminDao.getChangeStreamNamePrefix()); @@ -102,10 +99,7 @@ public class InitializeDoFnTest { Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = new InitializeDoFn( - daoFactory, - "app-profile", - startTime, - BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS); + daoFactory, startTime, BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS); initializeDoFn.processElement(outputReceiver); verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false)); } @@ -116,10 +110,7 @@ public class InitializeDoFnTest { Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = new InitializeDoFn( - daoFactory, - "app-profile", - startTime, - BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS); + daoFactory, startTime, BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS); initializeDoFn.processElement(outputReceiver); verify(outputReceiver, never()).output(any()); } @@ -139,51 +130,7 @@ public class InitializeDoFnTest { Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = new InitializeDoFn( - daoFactory, - "app-profile", - startTime, - BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS); - initializeDoFn.processElement(outputReceiver); - verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false)); - assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix())); - } - - @Test - public void testInitializeNewWithoutDNP() throws IOException { - // DNP row doesn't exist, so we output with provided start time and we clean up any existing - // rows under the prefix. - dataClient.mutateRow( - RowMutation.create( - tableId, - metadataTableAdminDao - .getChangeStreamNamePrefix() - .concat(ByteString.copyFromUtf8("existing_row"))) - .setCell( - MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123)); - Instant startTime = Instant.now(); - InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.NEW); - initializeDoFn.processElement(outputReceiver); - verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false)); - assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix())); - } - - @Test - public void testInitializeNewWithDNP() throws IOException { - metadataTableDao.updateDetectNewPartitionWatermark(Instant.now()); - dataClient.mutateRow( - RowMutation.create( - tableId, - metadataTableAdminDao - .getChangeStreamNamePrefix() - .concat(ByteString.copyFromUtf8("existing_row"))) - .setCell( - MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123)); - Instant startTime = Instant.now(); - InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.NEW); + daoFactory, startTime, BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS); initializeDoFn.processElement(outputReceiver); verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false)); assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix())); @@ -201,8 +148,7 @@ public class InitializeDoFnTest { MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123)); Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW); + new InitializeDoFn(daoFactory, startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW); initializeDoFn.processElement(outputReceiver); // We want to resume but there's no DNP row, so we resume from the startTime provided. verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false)); @@ -222,8 +168,7 @@ public class InitializeDoFnTest { MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123)); Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW); + new InitializeDoFn(daoFactory, startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW); initializeDoFn.processElement(outputReceiver); verify(outputReceiver, times(1)).output(new InitialPipelineState(resumeTime, true)); assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix())); @@ -241,8 +186,7 @@ public class InitializeDoFnTest { MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123)); Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, "app-profile", startTime, ExistingPipelineOptions.SKIP_CLEANUP); + new InitializeDoFn(daoFactory, startTime, ExistingPipelineOptions.SKIP_CLEANUP); initializeDoFn.processElement(outputReceiver); // Skip cleanup will always resume from startTime verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false)); @@ -264,8 +208,7 @@ public class InitializeDoFnTest { MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123)); Instant startTime = Instant.now(); InitializeDoFn initializeDoFn = - new InitializeDoFn( - daoFactory, "app-profile", startTime, ExistingPipelineOptions.SKIP_CLEANUP); + new InitializeDoFn(daoFactory, startTime, ExistingPipelineOptions.SKIP_CLEANUP); initializeDoFn.processElement(outputReceiver); // We don't want the pipeline to resume to avoid duplicates verify(outputReceiver, never()).output(any()); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java index cb82fc6d19c..b89b2bf15aa 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java @@ -20,7 +20,6 @@ package org.apache.beam.sdk.io.gcp.bigtable.changestreams.dofn; import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.TimestampConverter.toThreetenInstant; import static org.junit.Assert.assertEquals; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -94,7 +93,7 @@ public class ReadChangeStreamPartitionDoFnTest { sizeEstimator)) .thenReturn(readChangeStreamPartitionAction); - doFn = new ReadChangeStreamPartitionDoFn(heartbeatDuration, daoFactory, actionFactory, metrics); + doFn = new ReadChangeStreamPartitionDoFn(daoFactory, actionFactory, metrics); doFn.setSizeEstimator(sizeEstimator); } @@ -136,7 +135,7 @@ public class ReadChangeStreamPartitionDoFnTest { when(mockResponses.hasNext()).thenReturn(true, true, true); when(mockResponses.next()).thenReturn(mockMutation, mockMutation, mockMutation); when(mockStream.iterator()).thenReturn(mockResponses); - when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean())) + when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any())) .thenReturn(mockStream); when(watermarkEstimator.getState()).thenReturn(tenSecondsAgo);