Skip to content

Commit dc9179d

Browse files
scwhittlescwhittle
committed
address comments: ensure we handle other cases query terminates early. add tests.
1 parent def9bf9 commit dc9179d

2 files changed

Lines changed: 181 additions & 36 deletions

File tree

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

Lines changed: 30 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,6 @@ public class QueryChangeStreamAction {
168168
* @return a {@link ProcessContinuation#stop()} if a record timestamp could not be claimed or if
169169
* the partition processing has finished
170170
*/
171-
@SuppressWarnings("nullness")
172171
@VisibleForTesting
173172
public ProcessContinuation run(
174173
PartitionMetadata partition,
@@ -179,10 +178,9 @@ public ProcessContinuation run(
179178
final String token = partition.getPartitionToken();
180179
final Timestamp startTimestamp = tracker.currentRestriction().getFrom();
181180
final Timestamp endTimestamp = partition.getEndTimestamp();
181+
final boolean readToEndTimestamp = !endTimestamp.equals(MAX_INCLUSIVE_END_AT);
182182
final Timestamp changeStreamQueryEndTimestamp =
183-
endTimestamp.equals(MAX_INCLUSIVE_END_AT)
184-
? getNextReadChangeStreamEndTimestamp()
185-
: endTimestamp;
183+
readToEndTimestamp ? endTimestamp : getNextReadChangeStreamEndTimestamp();
186184

187185
// TODO: Potentially we can avoid this fetch, by enriching the runningAt timestamp when the
188186
// ReadChangeStreamPartitionDoFn#processElement is called
@@ -198,6 +196,7 @@ public ProcessContinuation run(
198196
RestrictionInterrupter<Timestamp> interrupter =
199197
RestrictionInterrupter.withSoftTimeout(RESTRICTION_TRACKER_TIMEOUT);
200198

199+
boolean stopAfterQuerySucceeds = readToEndTimestamp;
201200
try (ChangeStreamResultSet resultSet =
202201
changeStreamDao.changeStreamQuery(
203202
token, startTimestamp, changeStreamQueryEndTimestamp, partition.getHeartbeatMillis())) {
@@ -250,6 +249,9 @@ public ProcessContinuation run(
250249
tracker,
251250
interrupter,
252251
watermarkEstimator);
252+
// The PartitionEndRecord indicates that there are no more records expected
253+
// for this partition.
254+
stopAfterQuerySucceeds = true;
253255
} else if (record instanceof PartitionEventRecord) {
254256
maybeContinuation =
255257
partitionEventRecordAction.run(
@@ -272,27 +274,23 @@ public ProcessContinuation run(
272274
}
273275
}
274276
}
275-
bundleFinalizer.afterBundleCommit(
276-
Instant.now().plus(BUNDLE_FINALIZER_TIMEOUT),
277-
updateWatermarkCallback(token, watermarkEstimator));
278-
279277
} catch (SpannerException e) {
280278
/*
281279
If there is a split when a partition is supposed to be finished, the residual will try
282280
to perform a change stream query for an out of range interval. We ignore this error
283281
here, and the residual should be able to claim the end of the timestamp range, finishing
284282
the partition.
285283
*/
286-
if (isTimestampOutOfRange(e)) {
287-
LOG.info(
288-
"[{}] query change stream is out of range for {} to {}, finishing stream.",
289-
token,
290-
startTimestamp,
291-
endTimestamp,
292-
e);
293-
} else {
284+
if (!isTimestampOutOfRange(e)) {
294285
throw e;
295286
}
287+
LOG.info(
288+
"[{}] query change stream is out of range for {} to {}, finishing stream.",
289+
token,
290+
startTimestamp,
291+
endTimestamp,
292+
e);
293+
stopAfterQuerySucceeds = true;
296294
} catch (Exception e) {
297295
LOG.error(
298296
"[{}] query change stream had exception processing range {} to {}.",
@@ -305,22 +303,27 @@ public ProcessContinuation run(
305303

306304
LOG.debug(
307305
"[{}] change stream completed successfully up to {}", token, changeStreamQueryEndTimestamp);
308-
if (!tracker.tryClaim(changeStreamQueryEndTimestamp)) {
306+
Timestamp claimTimestamp =
307+
stopAfterQuerySucceeds ? endTimestamp : changeStreamQueryEndTimestamp;
308+
if (!tracker.tryClaim(claimTimestamp)) {
309309
return ProcessContinuation.stop();
310310
}
311+
bundleFinalizer.afterBundleCommit(
312+
Instant.now().plus(BUNDLE_FINALIZER_TIMEOUT),
313+
updateWatermarkCallback(token, watermarkEstimator));
311314

312-
if (changeStreamQueryEndTimestamp.equals(endTimestamp)) {
313-
LOG.debug("[{}] Finishing partition", token);
314-
// TODO: This should be performed after the commit succeeds. Since bundle finalizers are not
315-
// guaranteed to be called, this needs to be performed in a subsequent fused stage.
316-
partitionMetadataDao.updateToFinished(token);
317-
metrics.decActivePartitionReadCounter();
318-
LOG.info("[{}] After attempting to finish the partition", token);
319-
return ProcessContinuation.stop();
315+
if (!stopAfterQuerySucceeds) {
316+
LOG.debug("[{}] Rescheduling partition to resume reading", token);
317+
return ProcessContinuation.resume();
320318
}
321319

322-
LOG.info("[{}] Rescheduling partition where query completed due to not being finished", token);
323-
return ProcessContinuation.resume();
320+
LOG.debug("[{}] Finishing partition", token);
321+
// TODO: This should be performed after the commit succeeds. Since bundle finalizers are not
322+
// guaranteed to be called, this needs to be performed in a subsequent fused stage.
323+
partitionMetadataDao.updateToFinished(token);
324+
metrics.decActivePartitionReadCounter();
325+
LOG.info("[{}] After attempting to finish the partition", token);
326+
return ProcessContinuation.stop();
324327
}
325328

326329
private BundleFinalizer.Callback updateWatermarkCallback(

‎sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamActionTest.java‎

Lines changed: 151 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
import static org.mockito.Mockito.when;
3030

3131
import com.google.cloud.Timestamp;
32+
import com.google.cloud.spanner.ErrorCode;
33+
import com.google.cloud.spanner.SpannerExceptionFactory;
3234
import com.google.cloud.spanner.Struct;
3335
import java.util.Arrays;
3436
import java.util.Optional;
@@ -62,6 +64,7 @@
6264
public class QueryChangeStreamActionTest {
6365
private static final String PARTITION_TOKEN = "partitionToken";
6466
private static final Timestamp PARTITION_START_TIMESTAMP = Timestamp.ofTimeMicroseconds(10L);
67+
private static final Timestamp RECORD_TIMESTAMP = Timestamp.ofTimeMicroseconds(20L);
6568
private static final Timestamp PARTITION_END_TIMESTAMP = Timestamp.ofTimeMicroseconds(30L);
6669
private static final long PARTITION_HEARTBEAT_MILLIS = 30_000L;
6770
private static final Instant WATERMARK = Instant.now();
@@ -139,6 +142,22 @@ public void setUp() throws Exception {
139142
when(partitionMetadataMapper.from(row)).thenReturn(partition);
140143
}
141144

145+
void setupUnboundedPartition() {
146+
partition =
147+
PartitionMetadata.newBuilder()
148+
.setPartitionToken(PARTITION_TOKEN)
149+
.setParentTokens(Sets.newHashSet("parentToken"))
150+
.setStartTimestamp(PARTITION_START_TIMESTAMP)
151+
.setEndTimestamp(MAX_INCLUSIVE_END_AT)
152+
.setHeartbeatMillis(PARTITION_HEARTBEAT_MILLIS)
153+
.setState(SCHEDULED)
154+
.setWatermark(WATERMARK_TIMESTAMP)
155+
.setScheduledAt(Timestamp.now())
156+
.build();
157+
when(partitionMetadataMapper.from(any())).thenReturn(partition);
158+
when(restriction.getTo()).thenReturn(MAX_INCLUSIVE_END_AT);
159+
}
160+
142161
@Test
143162
public void testQueryChangeStreamWithDataChangeRecord() {
144163
final Struct rowAsStruct = mock(Struct.class);
@@ -148,7 +167,7 @@ public void testQueryChangeStreamWithDataChangeRecord() {
148167
final DataChangeRecord record1 = mock(DataChangeRecord.class);
149168
final DataChangeRecord record2 = mock(DataChangeRecord.class);
150169
when(record1.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
151-
when(record2.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
170+
when(record2.getRecordTimestamp()).thenReturn(RECORD_TIMESTAMP);
152171
when(changeStreamDao.changeStreamQuery(
153172
PARTITION_TOKEN,
154173
PARTITION_START_TIMESTAMP,
@@ -217,8 +236,8 @@ public void testQueryChangeStreamWithHeartbeatRecord() {
217236
final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
218237
final HeartbeatRecord record1 = mock(HeartbeatRecord.class);
219238
final HeartbeatRecord record2 = mock(HeartbeatRecord.class);
220-
when(record1.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
221-
when(record2.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
239+
when(record1.getRecordTimestamp()).thenReturn(RECORD_TIMESTAMP);
240+
when(record2.getRecordTimestamp()).thenReturn(RECORD_TIMESTAMP);
222241
when(changeStreamDao.changeStreamQuery(
223242
PARTITION_TOKEN,
224243
PARTITION_START_TIMESTAMP,
@@ -501,19 +520,19 @@ public void testQueryChangeStreamWithRestrictionFromAfterPartitionStartForPartit
501520
}
502521

503522
@Test
504-
public void testQueryChangeStreamWithPartitionEndRecord() {
523+
public void testQueryChangeStreamWithPartitionEndRecordBoundedRestriction() {
505524
final ChangeStreamResultSetMetadata resultSetMetadata =
506525
mock(ChangeStreamResultSetMetadata.class);
507526
final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
508527
final PartitionEndRecord record1 = mock(PartitionEndRecord.class);
509-
when(record1.getRecordTimestamp()).thenReturn(PARTITION_END_TIMESTAMP);
528+
when(record1.getRecordTimestamp()).thenReturn(RECORD_TIMESTAMP);
510529
when(changeStreamDao.changeStreamQuery(
511530
PARTITION_TOKEN,
512531
PARTITION_START_TIMESTAMP,
513532
PARTITION_END_TIMESTAMP,
514533
PARTITION_HEARTBEAT_MILLIS))
515534
.thenReturn(resultSet);
516-
when(resultSet.next()).thenReturn(true);
535+
when(resultSet.next()).thenReturn(true, false);
517536
when(resultSet.getMetadata()).thenReturn(resultSetMetadata);
518537
when(changeStreamRecordMapper.toChangeStreamRecords(partition, resultSet, resultSetMetadata))
519538
.thenReturn(Arrays.asList(record1));
@@ -523,8 +542,9 @@ public void testQueryChangeStreamWithPartitionEndRecord() {
523542
eq(restrictionTracker),
524543
any(RestrictionInterrupter.class),
525544
eq(watermarkEstimator)))
526-
.thenReturn(Optional.of(ProcessContinuation.stop()));
545+
.thenReturn(Optional.empty());
527546
when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);
547+
when(restrictionTracker.tryClaim(any(Timestamp.class))).thenReturn(true);
528548

529549
final ProcessContinuation result =
530550
action.run(
@@ -539,13 +559,66 @@ public void testQueryChangeStreamWithPartitionEndRecord() {
539559
any(RestrictionInterrupter.class),
540560
eq(watermarkEstimator));
541561
verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
562+
verify(restrictionTracker).tryClaim(PARTITION_END_TIMESTAMP);
563+
564+
verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any(), any());
565+
verify(heartbeatRecordAction, never()).run(any(), any(), any(), any(), any());
566+
verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any(), any());
567+
verify(partitionStartRecordAction, never()).run(any(), any(), any(), any(), any());
568+
verify(partitionEventRecordAction, never()).run(any(), any(), any(), any(), any());
569+
}
570+
571+
@Test
572+
public void testQueryChangeStreamWithPartitionEndRecordUnboundedRestriction() {
573+
setupUnboundedPartition();
574+
575+
final ChangeStreamResultSetMetadata resultSetMetadata =
576+
mock(ChangeStreamResultSetMetadata.class);
577+
final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
578+
final PartitionEndRecord record1 = mock(PartitionEndRecord.class);
579+
when(record1.getRecordTimestamp()).thenReturn(RECORD_TIMESTAMP);
580+
final ArgumentCaptor<Timestamp> timestampCaptor = ArgumentCaptor.forClass(Timestamp.class);
581+
when(changeStreamDao.changeStreamQuery(
582+
eq(PARTITION_TOKEN),
583+
eq(PARTITION_START_TIMESTAMP),
584+
timestampCaptor.capture(),
585+
eq(PARTITION_HEARTBEAT_MILLIS)))
586+
.thenReturn(resultSet);
587+
when(resultSet.next()).thenReturn(true, false);
588+
when(resultSet.getMetadata()).thenReturn(resultSetMetadata);
589+
when(changeStreamRecordMapper.toChangeStreamRecords(partition, resultSet, resultSetMetadata))
590+
.thenReturn(Arrays.asList(record1));
591+
when(partitionEndRecordAction.run(
592+
eq(partition),
593+
eq(record1),
594+
eq(restrictionTracker),
595+
any(RestrictionInterrupter.class),
596+
eq(watermarkEstimator)))
597+
.thenReturn(Optional.empty());
598+
when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);
599+
when(restrictionTracker.tryClaim(any(Timestamp.class))).thenReturn(true);
600+
601+
final ProcessContinuation result =
602+
action.run(
603+
partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
604+
605+
assertEquals(ProcessContinuation.stop(), result);
606+
assertNotEquals(MAX_INCLUSIVE_END_AT, timestampCaptor.getValue());
607+
verify(partitionEndRecordAction)
608+
.run(
609+
eq(partition),
610+
eq(record1),
611+
eq(restrictionTracker),
612+
any(RestrictionInterrupter.class),
613+
eq(watermarkEstimator));
614+
verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
615+
verify(restrictionTracker).tryClaim(MAX_INCLUSIVE_END_AT);
542616

543617
verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any(), any());
544618
verify(heartbeatRecordAction, never()).run(any(), any(), any(), any(), any());
545619
verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any(), any());
546620
verify(partitionStartRecordAction, never()).run(any(), any(), any(), any(), any());
547621
verify(partitionEventRecordAction, never()).run(any(), any(), any(), any(), any());
548-
verify(restrictionTracker, never()).tryClaim(any());
549622
}
550623

551624
@Test
@@ -614,7 +687,6 @@ public void testQueryChangeStreamWithStreamFinished() {
614687
partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
615688

616689
assertEquals(ProcessContinuation.stop(), result);
617-
verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
618690
verify(partitionMetadataDao).updateToFinished(PARTITION_TOKEN);
619691
verify(metrics).decActivePartitionReadCounter();
620692

@@ -672,6 +744,76 @@ public void testQueryChangeStreamFinishedWithResume() {
672744
verify(partitionEventRecordAction, never()).run(any(), any(), any(), any(), any());
673745
}
674746

747+
// Out of range indicates that we're beyond the end of the partition and should stop
748+
// processing.
749+
@Test
750+
public void testQueryChangeStreamWithOutOfRangeErrorOnUnboundedPartition() {
751+
setupUnboundedPartition();
752+
753+
final ArgumentCaptor<Timestamp> timestampCaptor = ArgumentCaptor.forClass(Timestamp.class);
754+
when(changeStreamDao.changeStreamQuery(
755+
eq(PARTITION_TOKEN),
756+
eq(PARTITION_START_TIMESTAMP),
757+
timestampCaptor.capture(),
758+
eq(PARTITION_HEARTBEAT_MILLIS)))
759+
.thenThrow(
760+
SpannerExceptionFactory.newSpannerException(
761+
ErrorCode.OUT_OF_RANGE, "Specified start_timestamp is invalid"));
762+
when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);
763+
when(restrictionTracker.tryClaim(any(Timestamp.class))).thenReturn(true);
764+
765+
final ProcessContinuation result =
766+
action.run(
767+
partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
768+
assertEquals(ProcessContinuation.stop(), result);
769+
assertNotEquals(MAX_INCLUSIVE_END_AT, timestampCaptor.getValue());
770+
771+
verify(restrictionTracker).tryClaim(MAX_INCLUSIVE_END_AT);
772+
verify(partitionMetadataDao).updateToFinished(PARTITION_TOKEN);
773+
verify(metrics).decActivePartitionReadCounter();
774+
verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
775+
776+
verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any(), any());
777+
verify(heartbeatRecordAction, never()).run(any(), any(), any(), any(), any());
778+
verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any(), any());
779+
verify(partitionStartRecordAction, never()).run(any(), any(), any(), any(), any());
780+
verify(partitionEndRecordAction, never()).run(any(), any(), any(), any(), any());
781+
verify(partitionEventRecordAction, never()).run(any(), any(), any(), any(), any());
782+
}
783+
784+
// Out of range indicates that we're beyond the end of the partition and should stop
785+
// processing.
786+
@Test
787+
public void testQueryChangeStreamWithOutOfRangeErrorOnBoundedPartition() {
788+
when(changeStreamDao.changeStreamQuery(
789+
eq(PARTITION_TOKEN),
790+
eq(PARTITION_START_TIMESTAMP),
791+
eq(PARTITION_END_TIMESTAMP),
792+
eq(PARTITION_HEARTBEAT_MILLIS)))
793+
.thenThrow(
794+
SpannerExceptionFactory.newSpannerException(
795+
ErrorCode.OUT_OF_RANGE, "Specified start_timestamp is invalid"));
796+
when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);
797+
when(restrictionTracker.tryClaim(any(Timestamp.class))).thenReturn(true);
798+
799+
final ProcessContinuation result =
800+
action.run(
801+
partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
802+
assertEquals(ProcessContinuation.stop(), result);
803+
804+
verify(restrictionTracker).tryClaim(PARTITION_END_TIMESTAMP);
805+
verify(partitionMetadataDao).updateToFinished(PARTITION_TOKEN);
806+
verify(metrics).decActivePartitionReadCounter();
807+
verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
808+
809+
verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any(), any());
810+
verify(heartbeatRecordAction, never()).run(any(), any(), any(), any(), any());
811+
verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any(), any());
812+
verify(partitionStartRecordAction, never()).run(any(), any(), any(), any(), any());
813+
verify(partitionEndRecordAction, never()).run(any(), any(), any(), any(), any());
814+
verify(partitionEventRecordAction, never()).run(any(), any(), any(), any(), any());
815+
}
816+
675817
private static class BundleFinalizerStub implements BundleFinalizer {
676818
@Override
677819
public void afterBundleCommit(Instant callbackExpiry, Callback callback) {

0 commit comments

Comments
 (0)