Skip to content

Commit 2c075b3

Browse files
jackdingilianjackdingilian
authored
Advance DetectNewPartition's watermark by aggregating the (#25906)
watermark of ReadChangeStreamPartitions
1 parent 8c90c64 commit 2c075b3

6 files changed

Lines changed: 467 additions & 60 deletions

File tree

‎sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/ByteStringRangeHelper.java‎

Lines changed: 100 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
2121
import com.google.protobuf.ByteString;
2222
import com.google.protobuf.TextFormat;
23+
import java.util.ArrayList;
24+
import java.util.Collections;
2325
import java.util.Comparator;
2426
import java.util.List;
2527
import java.util.stream.Collectors;
@@ -29,31 +31,6 @@
2931
/** Helper functions to evaluate the completeness of collection of ByteStringRanges. */
3032
@Internal
3133
public class ByteStringRangeHelper {
32-
/**
33-
* Returns formatted string of a partition for debugging.
34-
*
35-
* @param partition partition to format.
36-
* @return String representation of partition.
37-
*/
38-
public static String formatByteStringRange(ByteStringRange partition) {
39-
return "['"
40-
+ TextFormat.escapeBytes(partition.getStart())
41-
+ "','"
42-
+ TextFormat.escapeBytes(partition.getEnd())
43-
+ "')";
44-
}
45-
46-
/**
47-
* Convert partitions to a string for debugging.
48-
*
49-
* @param partitions to print
50-
* @return string representation of partitions
51-
*/
52-
public static String partitionsToString(List<ByteStringRange> partitions) {
53-
return partitions.stream()
54-
.map(ByteStringRangeHelper::formatByteStringRange)
55-
.collect(Collectors.joining(", ", "{", "}"));
56-
}
5734

5835
@VisibleForTesting
5936
static class PartitionComparator implements Comparator<ByteStringRange> {
@@ -83,6 +60,104 @@ public int compare(ByteStringRange first, ByteStringRange second) {
8360
}
8461
}
8562

63+
/**
64+
* Returns true if parentPartitions is a superset of childPartition.
65+
*
66+
* <p>If ordered parentPartitions row ranges form a contiguous range, and start key is before or
67+
* at childPartition's start key, and end key is at or after childPartition's end key, then
68+
* parentPartitions is a superset of childPartition.
69+
*
70+
* <p>Overlaps from parents are valid because arbitrary partitions can merge and they may overlap.
71+
* They will form a valid new partition. However, if there are any missing parent partitions, then
72+
* merge cannot happen with missing row ranges.
73+
*
74+
* @param parentPartitions list of partitions to determine if it forms a large contiguous range
75+
* @param childPartition the smaller partition
76+
* @return true if parentPartitions is a superset of childPartition, otherwise false.
77+
*/
78+
public static boolean isSuperset(
79+
List<ByteStringRange> parentPartitions, ByteStringRange childPartition) {
80+
// sort parentPartitions by starting key
81+
// iterate through, check open end key and close start key of each iteration to ensure no gaps.
82+
// first start key and last end key must be equal to or wider than child partition start and end
83+
// key.
84+
if (parentPartitions.isEmpty()) {
85+
return false;
86+
}
87+
parentPartitions.sort(new PartitionComparator());
88+
ByteString parentStartKey = parentPartitions.get(0).getStart();
89+
ByteString parentEndKey = parentPartitions.get(parentPartitions.size() - 1).getEnd();
90+
91+
return !childStartsBeforeParent(parentStartKey, childPartition.getStart())
92+
&& !childEndsAfterParent(parentEndKey, childPartition.getEnd())
93+
&& !gapsInParentPartitions(parentPartitions);
94+
}
95+
96+
/**
97+
* Convert partitions to a string for debugging.
98+
*
99+
* @param partitions to print
100+
* @return string representation of partitions
101+
*/
102+
public static String partitionsToString(List<ByteStringRange> partitions) {
103+
return partitions.stream()
104+
.map(ByteStringRangeHelper::formatByteStringRange)
105+
.collect(Collectors.joining(", ", "{", "}"));
106+
}
107+
108+
/**
109+
* Figure out if partitions cover the entire keyspace. If it doesn't, return a list of missing and
110+
* overlapping partitions.
111+
*
112+
* <p>partitions covers the entire key space if, when ordered, the end key is the same as the
113+
* start key of the next row range in the list, and the first start key is "" and the last end key
114+
* is "". There should be no overlap.
115+
*
116+
* @param partitions to determine if they cover entire keyspace
117+
* @return list of missing and overlapping partitions
118+
*/
119+
public static List<ByteStringRange> getMissingAndOverlappingPartitionsFromKeySpace(
120+
List<ByteStringRange> partitions) {
121+
if (partitions.isEmpty()) {
122+
return Collections.singletonList(ByteStringRange.create("", ""));
123+
}
124+
125+
List<ByteStringRange> missingPartitions = new ArrayList<>();
126+
127+
// sort partitions by start key
128+
// iterate through ensuring end key is lexicographically after next start key.
129+
partitions.sort(new PartitionComparator());
130+
131+
ByteString prevEnd = ByteString.EMPTY;
132+
for (ByteStringRange partition : partitions) {
133+
if (!partition.getStart().equals(prevEnd)) {
134+
ByteStringRange missingPartition = ByteStringRange.create(prevEnd, partition.getStart());
135+
missingPartitions.add(missingPartition);
136+
}
137+
prevEnd = partition.getEnd();
138+
}
139+
// Check that the last partition ends with "", otherwise it's missing.
140+
if (!prevEnd.equals(ByteString.EMPTY)) {
141+
ByteStringRange missingPartition = ByteStringRange.create(prevEnd, ByteString.EMPTY);
142+
missingPartitions.add(missingPartition);
143+
}
144+
return missingPartitions;
145+
}
146+
147+
/**
148+
* Returns formatted string of a partition for debugging.
149+
*
150+
* @param partition partition to format.
151+
* @return String representation of partition.
152+
*/
153+
public static String formatByteStringRange(ByteStringRange partition) {
154+
return "['"
155+
+ TextFormat.escapeBytes(partition.getStart())
156+
+ "','"
157+
+ TextFormat.escapeBytes(partition.getEnd())
158+
+ "')";
159+
}
160+
86161
private static boolean childStartsBeforeParent(
87162
ByteString parentStartKey, ByteString childStartKey) {
88163
// Check if the start key of the child partition comes before the start key of the entire
@@ -121,37 +196,4 @@ private static boolean gapsInParentPartitions(List<ByteStringRange> sortedParent
121196
}
122197
return false;
123198
}
124-
125-
/**
126-
* Returns true if parentPartitions is a superset of childPartition.
127-
*
128-
* <p>If ordered parentPartitions row ranges form a contiguous range, and start key is before or
129-
* at childPartition's start key, and end key is at or after childPartition's end key, then
130-
* parentPartitions is a superset of childPartition.
131-
*
132-
* <p>Overlaps from parents are valid because arbitrary partitions can merge and they may overlap.
133-
* They will form a valid new partition. However, if there are any missing parent partitions, then
134-
* merge cannot happen with missing row ranges.
135-
*
136-
* @param parentPartitions list of partitions to determine if it forms a large contiguous range
137-
* @param childPartition the smaller partition
138-
* @return true if parentPartitions is a superset of childPartition, otherwise false.
139-
*/
140-
public static boolean isSuperset(
141-
List<ByteStringRange> parentPartitions, ByteStringRange childPartition) {
142-
// sort parentPartitions by starting key
143-
// iterate through, check open end key and close start key of each iteration to ensure no gaps.
144-
// first start key and last end key must be equal to or wider than child partition start and end
145-
// key.
146-
if (parentPartitions.isEmpty()) {
147-
return false;
148-
}
149-
parentPartitions.sort(new PartitionComparator());
150-
ByteString parentStartKey = parentPartitions.get(0).getStart();
151-
ByteString parentEndKey = parentPartitions.get(parentPartitions.size() - 1).getEnd();
152-
153-
return !childStartsBeforeParent(parentStartKey, childPartition.getStart())
154-
&& !childEndsAfterParent(parentEndKey, childPartition.getEnd())
155-
&& !gapsInParentPartitions(parentPartitions);
156-
}
157199
}

‎sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/DetectNewPartitionsAction.java‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,23 @@
1717
*/
1818
package org.apache.beam.sdk.io.gcp.bigtable.changestreams.action;
1919

20+
import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.ByteStringRangeHelper.getMissingAndOverlappingPartitionsFromKeySpace;
21+
import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.ByteStringRangeHelper.partitionsToString;
22+
23+
import com.google.api.gax.rpc.ServerStream;
24+
import com.google.cloud.bigtable.data.v2.models.Range;
25+
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
26+
import com.google.cloud.bigtable.data.v2.models.Row;
2027
import com.google.protobuf.InvalidProtocolBufferException;
28+
import java.util.ArrayList;
29+
import java.util.HashMap;
30+
import java.util.List;
31+
import java.util.stream.Collectors;
2132
import org.apache.beam.sdk.annotations.Internal;
33+
import org.apache.beam.sdk.io.gcp.bigtable.changestreams.ByteStringRangeHelper;
2234
import org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamMetrics;
2335
import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.MetadataTableDao;
36+
import org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder;
2437
import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.PartitionRecord;
2538
import org.apache.beam.sdk.io.range.OffsetRange;
2639
import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer;
@@ -65,6 +78,83 @@ public DetectNewPartitionsAction(
6578
this.generateInitialPartitionsAction = generateInitialPartitionsAction;
6679
}
6780

81+
/**
82+
* Periodically advances DetectNewPartition's (DNP) watermark based on the watermark of all the
83+
* partitions recorded in the metadata table. We don't advance DNP's watermark on every run to
84+
* "now" because of a possible inconsistent state. DNP's watermark is used to hold the entire
85+
* pipeline's watermark back.
86+
*
87+
* <p>The low watermark is important to determine when to terminate the pipeline. During splits
88+
* and merges, the watermark step may appear to be higher than it actually is. If a partition, at
89+
* watermark 100, splits, it considered completed. If all other partitions including DNP have
90+
* watermark beyond 100, the low watermark of the pipeline is higher than 100. However, the split
91+
* partitions will have a watermark of 100 because they will resume from where the parent
92+
* partition has stopped. But this would mean the low watermark of the pipeline needs to move
93+
* backwards in time, which is not possible.
94+
*
95+
* <p>We "fix" this by using DNP's watermark to hold the pipeline's watermark down. DNP will
96+
* periodically scan the metadata table for all the partitions watermarks. It only advances its
97+
* watermark forward to the low watermark of the partitions. So in the case of a partition
98+
* split/merge the low watermark of the pipeline is held back by DNP. We guarantee correctness by
99+
* ensuring all the partitions exists in the metadata table in order to calculate the low
100+
* watermark. It is possible that some partitions might be missing in between split and merges.
101+
*
102+
* @param tracker restriction tracker to guide how frequently watermark should be advanced
103+
* @param watermarkEstimator watermark estimator to advance the watermark
104+
*/
105+
private void advanceWatermark(
106+
RestrictionTracker<OffsetRange, Long> tracker,
107+
ManualWatermarkEstimator<Instant> watermarkEstimator)
108+
throws InvalidProtocolBufferException {
109+
// We currently choose to update the watermark every 10 runs. We want to choose a number that is
110+
// frequent, so the watermark isn't lagged behind too far. Also not too frequent so we do not
111+
// overload the table with full table scans.
112+
if (tracker.currentRestriction().getFrom() % 10 == 0) {
113+
// Get partitions with a watermark set but skip rows w a lock and no watermark yet
114+
ServerStream<Row> rows = metadataTableDao.readFromMdTableStreamPartitionsWithWatermark();
115+
List<ByteStringRange> partitions = new ArrayList<>();
116+
HashMap<ByteStringRange, Instant> slowPartitions = new HashMap<>();
117+
Instant lowWatermark = Instant.ofEpochMilli(Long.MAX_VALUE);
118+
for (Row row : rows) {
119+
Instant watermark = MetadataTableEncoder.parseWatermarkFromRow(row);
120+
if (watermark == null) {
121+
continue;
122+
}
123+
// Update low watermark if watermark < low watermark.
124+
if (watermark.compareTo(lowWatermark) < 0) {
125+
lowWatermark = watermark;
126+
}
127+
Range.ByteStringRange partition =
128+
metadataTableDao.convertStreamPartitionRowKeyToPartition(row.getKey());
129+
partitions.add(partition);
130+
if (watermark.plus(DEBUG_WATERMARK_DELAY).isBeforeNow()) {
131+
slowPartitions.put(partition, watermark);
132+
}
133+
}
134+
List<Range.ByteStringRange> missingAndOverlappingPartitions =
135+
getMissingAndOverlappingPartitionsFromKeySpace(partitions);
136+
if (missingAndOverlappingPartitions.isEmpty()) {
137+
watermarkEstimator.setWatermark(lowWatermark);
138+
LOG.info("DNP: Updating watermark: " + watermarkEstimator.currentWatermark());
139+
} else {
140+
LOG.warn(
141+
"DNP: Could not update watermark because missing {}",
142+
partitionsToString(missingAndOverlappingPartitions));
143+
}
144+
if (!slowPartitions.isEmpty()) {
145+
LOG.warn(
146+
"DNP: Watermark is being held back by the following partitions: {}",
147+
slowPartitions.entrySet().stream()
148+
.map(
149+
e ->
150+
ByteStringRangeHelper.formatByteStringRange(e.getKey())
151+
+ " => "
152+
+ e.getValue())
153+
.collect(Collectors.joining(", ", "{", "}")));
154+
}
155+
}
156+
}
157+
68158
/**
69159
* Perform the necessary steps to manage initial set of partitions and new partitions. Currently,
70160
* we set to process new partitions every second.
@@ -102,6 +192,8 @@ public ProcessContinuation run(
102192
return ProcessContinuation.stop();
103193
}
104194

195+
advanceWatermark(tracker, watermarkEstimator);
196+
105197
return ProcessContinuation.resume().withResumeDelay(Duration.standardSeconds(1));
106198
}
107199
}

‎sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java‎

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,21 @@
1717
*/
1818
package org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao;
1919

20+
import static com.google.cloud.bigtable.data.v2.models.Filters.FILTERS;
2021
import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.MetadataTableAdminDao.DETECT_NEW_PARTITION_SUFFIX;
2122
import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.MetadataTableAdminDao.NEW_PARTITION_PREFIX;
2223
import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.MetadataTableAdminDao.STREAM_PARTITION_PREFIX;
2324

2425
import com.google.api.gax.rpc.ServerStream;
2526
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
2627
import com.google.cloud.bigtable.data.v2.models.ChangeStreamContinuationToken;
27-
import com.google.cloud.bigtable.data.v2.models.Filters;
2828
import com.google.cloud.bigtable.data.v2.models.Query;
2929
import com.google.cloud.bigtable.data.v2.models.Range;
30+
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
3031
import com.google.cloud.bigtable.data.v2.models.Row;
3132
import com.google.cloud.bigtable.data.v2.models.RowMutation;
3233
import com.google.protobuf.ByteString;
34+
import com.google.protobuf.InvalidProtocolBufferException;
3335
import javax.annotation.Nullable;
3436
import org.apache.beam.sdk.annotations.Internal;
3537
import org.joda.time.Instant;
@@ -90,6 +92,21 @@ private ByteString getFullDetectNewPartition() {
9092
return changeStreamNamePrefix.concat(DETECT_NEW_PARTITION_SUFFIX);
9193
}
9294

95+
/**
96+
* Convert stream partition row key to partition to process metadata read from Bigtable.
97+
*
98+
* <p>RowKey should be directly from Cloud Bigtable and not altered in any way.
99+
*
100+
* @param rowKey row key from Cloud Bigtable
101+
* @return partition extracted from rowKey
102+
* @throws InvalidProtocolBufferException if conversion from rowKey to partition fails
103+
*/
104+
public ByteStringRange convertStreamPartitionRowKeyToPartition(ByteString rowKey)
105+
throws InvalidProtocolBufferException {
106+
int prefixLength = changeStreamNamePrefix.size() + STREAM_PARTITION_PREFIX.size();
107+
return ByteStringRange.toByteStringRange(rowKey.substring(prefixLength));
108+
}
109+
93110
/**
94111
* Convert partition to a Stream Partition row key to query for metadata of partitions that are
95112
* currently being streamed.
@@ -125,7 +142,7 @@ public ServerStream<Row> readNewPartitions() {
125142
Query query =
126143
Query.create(tableId)
127144
.prefix(getFullNewPartitionPrefix())
128-
.filter(Filters.FILTERS.limit().cellsPerColumn(1));
145+
.filter(FILTERS.limit().cellsPerColumn(1));
129146
return dataClient.readRows(query);
130147
}
131148

@@ -175,6 +192,25 @@ private void writeNewPartition(
175192
dataClient.mutateRow(rowMutation);
176193
}
177194

195+
/**
196+
* @return stream of partitions currently being streamed by the beam job that have set a
197+
* watermark.
198+
*/
199+
public ServerStream<Row> readFromMdTableStreamPartitionsWithWatermark() {
200+
// We limit to the latest value per column.
201+
Query query =
202+
Query.create(tableId)
203+
.prefix(getFullStreamPartitionPrefix())
204+
.filter(
205+
FILTERS
206+
.chain()
207+
.filter(FILTERS.limit().cellsPerColumn(1))
208+
.filter(FILTERS.family().exactMatch(MetadataTableAdminDao.CF_WATERMARK))
209+
.filter(
210+
FILTERS.qualifier().exactMatch(MetadataTableAdminDao.QUALIFIER_DEFAULT)));
211+
return dataClient.readRows(query);
212+
}
213+
178214
/**
179215
* Update the metadata for the rowKey. This helper adds necessary prefixes to the row key.
180216
*

0 commit comments

Comments
 (0)