This is an automated email from the ASF dual-hosted git repository.
davidzollo pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/seatunnel.git
The following commit(s) were added to refs/heads/dev by this push:
new 244455533b [Improve][Connector-V2] Harden Google Bigtable source-side
state synchronization (#11144)
244455533b is described below
commit 244455533b2a25be916bc8a9ad5e3557e6668de4
Author: 雷炯 <[email protected]>
AuthorDate: Tue Jul 21 22:41:17 2026 +0800
[Improve][Connector-V2] Harden Google Bigtable source-side state
synchronization (#11144)
Co-authored-by: Cursor <[email protected]>
---
.../bigtable/source/BigtableSourceReader.java | 45 +-
.../source/BigtableSourceSplitEnumerator.java | 52 ++-
.../bigtable/source/BigtableSourceState.java | 25 +-
.../source/BigtableSourceSplitEnumeratorTest.java | 152 +++++++
.../source/BigtableSourceStateRecoveryTest.java | 458 +++++++++++++++++++++
5 files changed, 701 insertions(+), 31 deletions(-)
diff --git
a/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceReader.java
b/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceReader.java
index 3dcfa310b2..fecec582d3 100644
---
a/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceReader.java
+++
b/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceReader.java
@@ -64,6 +64,13 @@ public class BigtableSourceReader implements
SourceReader<SeaTunnelRow, Bigtable
/** The split currently being read. Persisted in checkpoint to survive
failover. */
private volatile BigtableSourceSplit currentSplit = null;
+ /**
+ * Guards compound reader-state transitions so a concurrent {@link
#snapshotState(long)} never
+ * observes a half-updated view of {@code pendingSplits} and {@code
currentSplit}. The actual
+ * {@code readRows} network scan is intentionally kept outside this lock.
+ */
+ private final Object stateLock = new Object();
+
/**
* Set of field names that map to the Bigtable row key. Populated from
{@code rowkey_column}
* config; falls back to the literal field name {@value ROW_KEY_FIELD}
when not configured.
@@ -109,13 +116,23 @@ public class BigtableSourceReader implements
SourceReader<SeaTunnelRow, Bigtable
@Override
public void pollNext(Collector<SeaTunnelRow> output) throws Exception {
- final BigtableSourceSplit split = pendingSplits.poll();
+ final BigtableSourceSplit split;
+ synchronized (stateLock) {
+ // Poll and assign currentSplit atomically so a concurrent
snapshotState() always sees a
+ // consistent (pendingSplits, currentSplit) pair, and the
in-flight split can be
+ // re-enqueued on restore.
+ split = pendingSplits.poll();
+ if (Objects.nonNull(split)) {
+ currentSplit = split;
+ }
+ }
if (Objects.nonNull(split)) {
- // Assign currentSplit before reading so checkpoints taken during
readSplit()
- // include it and can re-enqueue it on restore.
- currentSplit = split;
+ // The network scan runs outside stateLock; only output.collect()
takes the checkpoint
+ // lock inside readSplit().
readSplit(split, output);
- currentSplit = null;
+ synchronized (stateLock) {
+ currentSplit = null;
+ }
} else if (noMoreSplits && pendingSplits.isEmpty()) {
log.info("Closed the bounded Bigtable source");
context.signalNoMoreElement();
@@ -234,18 +251,22 @@ public class BigtableSourceReader implements
SourceReader<SeaTunnelRow, Bigtable
@Override
public List<BigtableSourceSplit> snapshotState(long checkpointId) {
- List<BigtableSourceSplit> state = new ArrayList<>();
- // Include the split currently being read so it can be re-enqueued on
restore.
- if (currentSplit != null) {
- state.add(currentSplit);
+ synchronized (stateLock) {
+ List<BigtableSourceSplit> state = new ArrayList<>();
+ // Include the split currently being read so it can be re-enqueued
on restore.
+ if (currentSplit != null) {
+ state.add(currentSplit);
+ }
+ state.addAll(pendingSplits);
+ return state;
}
- state.addAll(pendingSplits);
- return state;
}
@Override
public void addSplits(List<BigtableSourceSplit> splits) {
- pendingSplits.addAll(splits);
+ synchronized (stateLock) {
+ pendingSplits.addAll(splits);
+ }
}
@Override
diff --git
a/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumerator.java
index 3d41a2ad39..df72fda91d 100644
---
a/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumerator.java
@@ -47,31 +47,43 @@ public class BigtableSourceSplitEnumerator
private Set<BigtableSourceSplit> pendingSplits;
private boolean initialized = false;
+ /**
+ * Guards the shared assignment state ({@code assignedSplits}, {@code
pendingSplits}, {@code
+ * initialized}) against concurrent enumerator callbacks. {@link
#initializePendingSplits()} and
+ * {@link #assignSplit(int)} must only run while this lock is held.
+ */
+ private final Object stateLock = new Object();
+
public BigtableSourceSplitEnumerator(
Context<BigtableSourceSplit> context, BigtableParameters
parameters) {
- this(context, parameters, new HashSet<>());
+ this(context, parameters, null);
}
public BigtableSourceSplitEnumerator(
Context<BigtableSourceSplit> context,
BigtableParameters parameters,
BigtableSourceState sourceState) {
- this(context, parameters, sourceState.getAssignedSplits());
- }
-
- private BigtableSourceSplitEnumerator(
- Context<BigtableSourceSplit> context,
- BigtableParameters parameters,
- Set<BigtableSourceSplit> assignedSplits) {
this.context = context;
this.parameters = parameters;
- this.assignedSplits = new HashSet<>(assignedSplits);
+ if (sourceState == null) {
+ this.assignedSplits = new HashSet<>();
+ this.pendingSplits = new HashSet<>();
+ this.initialized = false;
+ } else {
+ this.assignedSplits = new
HashSet<>(sourceState.getAssignedSplits());
+ // Restore pending splits first so a
returned-but-not-yet-reassigned split survives
+ // recovery even when its ID is already present in assignedSplits.
+ this.pendingSplits = new HashSet<>(sourceState.getPendingSplits());
+ // Only skip table-split discovery when the checkpoint captured
real enumerator
+ // progress.
+ // An empty-empty checkpoint (before any reader registered) must
still discover splits.
+ this.initialized = !this.assignedSplits.isEmpty() ||
!this.pendingSplits.isEmpty();
+ }
}
@Override
public void open() {
- this.pendingSplits = new HashSet<>();
- this.initialized = false;
+ // State is fully initialized in the constructor; nothing to reset on
open().
}
@Override
@@ -87,9 +99,11 @@ public class BigtableSourceSplitEnumerator
@Override
public void addSplitsBack(List<BigtableSourceSplit> splits, int subtaskId)
{
if (!splits.isEmpty()) {
- pendingSplits.addAll(splits);
- if (context.registeredReaders().contains(subtaskId)) {
- assignSplit(subtaskId);
+ synchronized (stateLock) {
+ pendingSplits.addAll(splits);
+ if (context.registeredReaders().contains(subtaskId)) {
+ assignSplit(subtaskId);
+ }
}
}
}
@@ -101,13 +115,17 @@ public class BigtableSourceSplitEnumerator
@Override
public void registerReader(int subtaskId) {
- initializePendingSplits();
- assignSplit(subtaskId);
+ synchronized (stateLock) {
+ initializePendingSplits();
+ assignSplit(subtaskId);
+ }
}
@Override
public BigtableSourceState snapshotState(long checkpointId) throws
Exception {
- return new BigtableSourceState(assignedSplits);
+ synchronized (stateLock) {
+ return new BigtableSourceState(assignedSplits, pendingSplits);
+ }
}
@Override
diff --git
a/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceState.java
b/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceState.java
index e44fd024a9..932a54503c 100644
---
a/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceState.java
+++
b/seatunnel-connectors-v2/connector-google-bigtable/src/main/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceState.java
@@ -18,6 +18,8 @@
package org.apache.seatunnel.connectors.seatunnel.bigtable.source;
import java.io.Serializable;
+import java.util.Collections;
+import java.util.HashSet;
import java.util.Set;
public class BigtableSourceState implements Serializable {
@@ -25,12 +27,31 @@ public class BigtableSourceState implements Serializable {
private static final long serialVersionUID = 1L;
private final Set<BigtableSourceSplit> assignedSplits;
+ /**
+ * Splits returned via {@code addSplitsBack()} but not yet reassigned.
Added in #11144; absent
+ * (null) when an older checkpoint payload is deserialized with the same
{@link
+ * #serialVersionUID}.
+ */
+ private final Set<BigtableSourceSplit> pendingSplits;
- public BigtableSourceState(Set<BigtableSourceSplit> assignedSplits) {
- this.assignedSplits = assignedSplits;
+ public BigtableSourceState(
+ Set<BigtableSourceSplit> assignedSplits, Set<BigtableSourceSplit>
pendingSplits) {
+ // Defensive copies: never alias the enumerator's live mutable sets
into checkpoint state,
+ // otherwise later mutations could corrupt an already-snapshotted
state object.
+ this.assignedSplits = copyOrEmpty(assignedSplits);
+ this.pendingSplits = copyOrEmpty(pendingSplits);
}
public Set<BigtableSourceSplit> getAssignedSplits() {
return assignedSplits;
}
+
+ public Set<BigtableSourceSplit> getPendingSplits() {
+ // Null only when Java deserializes a pre-#11144 checkpoint that
lacked this field.
+ return pendingSplits == null ? Collections.emptySet() : pendingSplits;
+ }
+
+ private static Set<BigtableSourceSplit>
copyOrEmpty(Set<BigtableSourceSplit> splits) {
+ return splits == null ? new HashSet<>() : new HashSet<>(splits);
+ }
}
diff --git
a/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumeratorTest.java
index b9d764b34a..b0767c7252 100644
---
a/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumeratorTest.java
+++
b/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceSplitEnumeratorTest.java
@@ -17,9 +17,24 @@
package org.apache.seatunnel.connectors.seatunnel.bigtable.source;
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import
org.apache.seatunnel.connectors.seatunnel.bigtable.config.BigtableParameters;
+
import org.junit.jupiter.api.Test;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
public class BigtableSourceSplitEnumeratorTest {
@@ -37,4 +52,141 @@ public class BigtableSourceSplitEnumeratorTest {
assertEquals("aaa", split.getStartRowKey());
assertEquals("zzz", split.getEndRowKey());
}
+
+ /**
+ * Regression for the returned-split recovery window: a split requeued via
{@link
+ * BigtableSourceSplitEnumerator#addSplitsBack} but not yet reassigned
must survive enumerator
+ * checkpoint / restore and be reassigned to the reader.
+ */
+ @Test
+ void testReturnedSplitSurvivesCheckpointRestoreAndReassign() throws
Exception {
+ TestingContext context = new TestingContext(1);
+ BigtableParameters parameters =
+ BigtableParameters.builder()
+ .projectId("test-project")
+ .instanceId("test-instance")
+ .table("test-table")
+ .build();
+
+ BigtableSourceSplitEnumerator enumerator =
+ new BigtableSourceSplitEnumerator(context, parameters);
+ enumerator.open();
+
+ context.registerReaderForTest(0);
+ enumerator.registerReader(0);
+ BigtableSourceSplit split = context.getLastAssignedSplit(0);
+ assertEquals("bigtable_source_split_0", split.splitId());
+
+ // Simulate failover: reader returns its split before the enumerator
can reassign it.
+ context.unregisterReaderForTest(0);
+ enumerator.addSplitsBack(Collections.singletonList(split), 0);
+
+ BigtableSourceState checkpoint = enumerator.snapshotState(1L);
+ assertTrue(checkpoint.getPendingSplits().contains(split));
+ assertTrue(checkpoint.getAssignedSplits().contains(split));
+
+ BigtableSourceSplitEnumerator restored =
+ new BigtableSourceSplitEnumerator(context, parameters,
checkpoint);
+ restored.open();
+
+ context.clearAssignments();
+ context.registerReaderForTest(0);
+ restored.registerReader(0);
+
+ assertEquals(1, context.getAssignedSplitCount(0));
+ assertEquals(split.splitId(),
context.getLastAssignedSplit(0).splitId());
+ assertEquals(0, restored.currentUnassignedSplitSize());
+ }
+
+ @Test
+ void testSnapshotStateDefensivelyCopiesPendingSplits() throws Exception {
+ TestingContext context = new TestingContext(1);
+ BigtableParameters parameters =
+ BigtableParameters.builder()
+ .projectId("test-project")
+ .instanceId("test-instance")
+ .table("test-table")
+ .build();
+
+ BigtableSourceSplitEnumerator enumerator =
+ new BigtableSourceSplitEnumerator(context, parameters);
+ enumerator.open();
+
+ context.registerReaderForTest(0);
+ enumerator.registerReader(0);
+ BigtableSourceSplit split = context.getLastAssignedSplit(0);
+
+ context.unregisterReaderForTest(0);
+ enumerator.addSplitsBack(Collections.singletonList(split), 0);
+
+ BigtableSourceState checkpoint = enumerator.snapshotState(1L);
+ assertEquals(1, checkpoint.getPendingSplits().size());
+ checkpoint.getPendingSplits().clear();
+
+ assertEquals(1, enumerator.currentUnassignedSplitSize());
+ }
+
+ private static class TestingContext
+ implements SourceSplitEnumerator.Context<BigtableSourceSplit> {
+
+ private final int parallelism;
+ private final Set<Integer> registeredReaders = new HashSet<>();
+ private final Map<Integer, List<BigtableSourceSplit>> assignments =
new HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ }
+
+ void registerReaderForTest(int subtaskId) {
+ registeredReaders.add(subtaskId);
+ }
+
+ void unregisterReaderForTest(int subtaskId) {
+ registeredReaders.remove(subtaskId);
+ }
+
+ void clearAssignments() {
+ assignments.clear();
+ }
+
+ int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+
+ BigtableSourceSplit getLastAssignedSplit(int subtaskId) {
+ List<BigtableSourceSplit> splits = assignments.get(subtaskId);
+ return splits.get(splits.size() - 1);
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return registeredReaders;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<BigtableSourceSplit>
splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+ }
}
diff --git
a/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceStateRecoveryTest.java
b/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceStateRecoveryTest.java
new file mode 100644
index 0000000000..e67b0b4ec5
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-google-bigtable/src/test/java/org/apache/seatunnel/connectors/seatunnel/bigtable/source/BigtableSourceStateRecoveryTest.java
@@ -0,0 +1,458 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.bigtable.source;
+
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.Collector;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceReader;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import org.apache.seatunnel.api.table.type.BasicType;
+import org.apache.seatunnel.api.table.type.SeaTunnelRow;
+import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
+import
org.apache.seatunnel.connectors.seatunnel.bigtable.client.BigtableClient;
+import
org.apache.seatunnel.connectors.seatunnel.bigtable.config.BigtableParameters;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import com.google.api.gax.rpc.ServerStream;
+import com.google.cloud.bigtable.data.v2.BigtableDataClient;
+import com.google.cloud.bigtable.data.v2.models.Query;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.ArrayList;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Deep regression tests for Bigtable source checkpoint / recovery semantics
introduced in #11144.
+ */
+class BigtableSourceStateRecoveryTest {
+
+ /**
+ * Bytes from the pre-#11144 {@code BigtableSourceState} (assignedSplits
only, same
+ * serialVersionUID) for split {@code bigtable_source_split_0} with row
range [start, end).
+ */
+ private static final String LEGACY_ASSIGNED_ONLY_STATE_BASE64 =
+
"rO0ABXNyAE1vcmcuYXBhY2hlLnNlYXR1bm5lbC5jb25uZWN0b3JzLnNlYXR1bm5lbC5iaWd0YWJsZS5zb3VyY2UuQmlndGFibGVTb3VyY2VTdGF0ZQAAAAAAAAABAgABTAAOYXNzaWduZWRTcGxpdHN0AA9MamF2YS91dGlsL1NldDt4cHNyABFqYXZhLnV0aWwuSGFzaFNldLpEhZWWuLc0AwAAeHB3DAAAABA/QAAAAAAAAXNyAE1vcmcuYXBhY2hlLnNlYXR1bm5lbC5jb25uZWN0b3JzLnNlYXR1bm5lbC5iaWd0YWJsZS5zb3VyY2UuQmlndGFibGVTb3VyY2VTcGxpdAAAAAAAAAABAgADTAAJZW5kUm93S2V5dAASTGphdmEvbGFuZy9TdHJpbmc7TAAHc3BsaXRJZHEAfgAGTAALc3RhcnRSb3dLZXlxAH4ABnhwdAADZW5kdAAXYmlndGFibGV
[...]
+
+ private static final BigtableParameters PARAMETERS =
+ BigtableParameters.builder()
+ .projectId("test-project")
+ .instanceId("test-instance")
+ .table("test-table")
+ .build();
+
+ /**
+ * If the enumerator checkpoints before any reader registers, both sets
are empty. Restore must
+ * still discover and assign the table split on the next {@link
+ * BigtableSourceSplitEnumerator#registerReader(int)}.
+ */
+ @Test
+ void testEmptyEnumeratorCheckpointStillDiscoversSplitsOnRestore() throws
Exception {
+ TestingContext context = new TestingContext(1);
+ BigtableSourceSplitEnumerator enumerator =
+ new BigtableSourceSplitEnumerator(context, PARAMETERS);
+ enumerator.open();
+
+ BigtableSourceState emptyCheckpoint = enumerator.snapshotState(1L);
+ assertTrue(emptyCheckpoint.getAssignedSplits().isEmpty());
+ assertTrue(emptyCheckpoint.getPendingSplits().isEmpty());
+
+ BigtableSourceSplitEnumerator restored =
+ new BigtableSourceSplitEnumerator(context, PARAMETERS,
emptyCheckpoint);
+ restored.open();
+
+ context.registerReaderForTest(0);
+ restored.registerReader(0);
+
+ assertEquals(1, context.getAssignedSplitCount(0));
+ assertEquals("bigtable_source_split_0",
context.getLastAssignedSplit(0).splitId());
+ }
+
+ /**
+ * Documents Daniel's blocker: without pendingSplits in checkpoint state,
a returned split whose
+ * ID is already in assignedSplits is dropped after restore.
+ */
+ @Test
+ void testReturnedSplitDroppedWhenCheckpointOmitsPending() throws Exception
{
+ TestingContext context = new TestingContext(1);
+ BigtableSourceSplitEnumerator enumerator =
+ new BigtableSourceSplitEnumerator(context, PARAMETERS);
+ enumerator.open();
+
+ context.registerReaderForTest(0);
+ enumerator.registerReader(0);
+ BigtableSourceSplit split = context.getLastAssignedSplit(0);
+
+ context.unregisterReaderForTest(0);
+ enumerator.addSplitsBack(Collections.singletonList(split), 0);
+ assertEquals(1, enumerator.currentUnassignedSplitSize());
+
+ // Simulate pre-fix snapshot: only assignedSplits, pending lost.
+ BigtableSourceState buggyCheckpoint =
+ new BigtableSourceState(
+ new HashSet<>(Collections.singleton(split)),
Collections.emptySet());
+
+ TestingContext restoreContext = new TestingContext(1);
+ BigtableSourceSplitEnumerator restored =
+ new BigtableSourceSplitEnumerator(restoreContext, PARAMETERS,
buggyCheckpoint);
+ restored.open();
+
+ restoreContext.registerReaderForTest(0);
+ restored.registerReader(0);
+
+ assertEquals(
+ 0,
+ restoreContext.getAssignedSplitCount(0),
+ "Without pendingSplits in checkpoint, returned split is
dropped on restore");
+ }
+
+ /**
+ * Checkpoints written before #11144 only serialized assignedSplits. After
Java deserialization
+ * into the new class shape, pendingSplits is null; restore must not fail.
+ */
+ @Test
+ void testLegacyCheckpointWithNullPendingSplitsRestoresEnumerator() throws
Exception {
+ BigtableSourceSplit split = new BigtableSourceSplit(0, "start", "end");
+ BigtableSourceState legacyState =
+ new BigtableSourceState(
+ new HashSet<>(Collections.singleton(split)),
Collections.emptySet());
+ java.lang.reflect.Field pendingField =
+ BigtableSourceState.class.getDeclaredField("pendingSplits");
+ pendingField.setAccessible(true);
+ pendingField.set(legacyState, null);
+
+ assertTrue(legacyState.getPendingSplits().isEmpty());
+
+ TestingContext context = new TestingContext(1);
+ BigtableSourceSplitEnumerator restored =
+ new BigtableSourceSplitEnumerator(context, PARAMETERS,
legacyState);
+ restored.open();
+
+ context.registerReaderForTest(0);
+ restored.registerReader(0);
+
+ assertEquals(0, restored.currentUnassignedSplitSize());
+ }
+
+ /**
+ * Round-trip bytes produced by the pre-#11144 single-field {@code
BigtableSourceState} shape
+ * (same {@code serialVersionUID}) must deserialize into the new class
without restore failure.
+ */
+ @Test
+ void testDeserializePrePatchBigtableSourceStateBytes() throws Exception {
+ byte[] legacyBytes =
Base64.getDecoder().decode(LEGACY_ASSIGNED_ONLY_STATE_BASE64);
+
+ BigtableSourceState deserialized;
+ try (ObjectInputStream ois = new ObjectInputStream(new
ByteArrayInputStream(legacyBytes))) {
+ deserialized = (BigtableSourceState) ois.readObject();
+ }
+
+ assertEquals(1, deserialized.getAssignedSplits().size());
+ assertTrue(deserialized.getPendingSplits().isEmpty());
+
+ TestingContext context = new TestingContext(1);
+ BigtableSourceSplitEnumerator restored =
+ new BigtableSourceSplitEnumerator(context, PARAMETERS,
deserialized);
+ restored.open();
+
+ context.registerReaderForTest(0);
+ restored.registerReader(0);
+
+ assertEquals(0, restored.currentUnassignedSplitSize());
+ }
+
+ /**
+ * End-to-end reader + enumerator failover: in-flight reader split flows
back via addSplitsBack.
+ */
+ @Test
+ void testReaderEnumeratorFailoverHandoff() throws Exception {
+ TestingContext enumContext = new TestingContext(1);
+ BigtableSourceSplitEnumerator enumerator =
+ new BigtableSourceSplitEnumerator(enumContext, PARAMETERS);
+ enumerator.open();
+ enumContext.registerReaderForTest(0);
+ enumerator.registerReader(0);
+ BigtableSourceSplit enumeratorSplit =
enumContext.getLastAssignedSplit(0);
+
+ BigtableSourceReader reader = createReaderWithMockedEmptyStream();
+ reader.addSplits(Collections.singletonList(enumeratorSplit));
+ reader.handleNoMoreSplits();
+
+ final AtomicReference<List<BigtableSourceSplit>> readerCheckpoint =
new AtomicReference<>();
+ mockEmptyReadStream(reader, () ->
readerCheckpoint.set(reader.snapshotState(1L)));
+
+ Collector<SeaTunnelRow> collector = mockCollector();
+ reader.pollNext(collector);
+
+ List<BigtableSourceSplit> returnedSplits = readerCheckpoint.get();
+ assertNotNull(returnedSplits);
+ assertEquals(1, returnedSplits.size());
+
+ enumContext.unregisterReaderForTest(0);
+ enumerator.addSplitsBack(returnedSplits, 0);
+
+ BigtableSourceState enumCheckpoint = enumerator.snapshotState(2L);
+
assertTrue(enumCheckpoint.getPendingSplits().contains(enumeratorSplit));
+
+ TestingContext restoreContext = new TestingContext(1);
+ BigtableSourceSplitEnumerator restoredEnum =
+ new BigtableSourceSplitEnumerator(restoreContext, PARAMETERS,
enumCheckpoint);
+ restoredEnum.open();
+ restoreContext.registerReaderForTest(0);
+ restoredEnum.registerReader(0);
+
+ assertEquals(1, restoreContext.getAssignedSplitCount(0));
+ assertEquals(enumeratorSplit.splitId(),
restoreContext.getLastAssignedSplit(0).splitId());
+ }
+
+ /** Concurrent reader snapshot and addSplits must remain thread-safe (no
torn exceptions). */
+ @Test
+ void testConcurrentReaderSnapshotAndAddSplits() throws Exception {
+ BigtableSourceReader reader = createReaderWithMockedEmptyStream();
+ BigtableSourceSplit splitA = new BigtableSourceSplit(0, "a", "m");
+ BigtableSourceSplit splitB = new BigtableSourceSplit(1, "m", "z");
+
+ ExecutorService pool = Executors.newFixedThreadPool(4);
+ CountDownLatch start = new CountDownLatch(1);
+ List<List<BigtableSourceSplit>> snapshots =
Collections.synchronizedList(new ArrayList<>());
+
+ for (int i = 0; i < 100; i++) {
+ pool.submit(
+ () -> {
+ try {
+ start.await();
+ snapshots.add(reader.snapshotState(1L));
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ pool.submit(
+ () -> {
+ try {
+ start.await();
+
reader.addSplits(Collections.singletonList(splitA));
+
reader.addSplits(Collections.singletonList(splitB));
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ }
+
+ start.countDown();
+ pool.shutdown();
+ assertTrue(pool.awaitTermination(10, TimeUnit.SECONDS));
+ assertEquals(100, snapshots.size());
+ for (List<BigtableSourceSplit> snapshot : snapshots) {
+ assertNotNull(snapshot);
+ }
+ }
+
+ @Test
+ void testBigtableSourceStateJavaSerializationRoundTrip() throws Exception {
+ BigtableSourceSplit split = new BigtableSourceSplit(0, "start", "end");
+ BigtableSourceState original =
+ new BigtableSourceState(
+ new HashSet<>(Collections.singleton(split)),
+ new HashSet<>(Collections.singleton(split)));
+
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ try (ObjectOutputStream oos = new ObjectOutputStream(bos)) {
+ oos.writeObject(original);
+ }
+
+ BigtableSourceState restored;
+ try (ObjectInputStream ois =
+ new ObjectInputStream(new
ByteArrayInputStream(bos.toByteArray()))) {
+ restored = (BigtableSourceState) ois.readObject();
+ }
+
+ assertEquals(1, restored.getAssignedSplits().size());
+ assertEquals(1, restored.getPendingSplits().size());
+ assertEquals(
+ "bigtable_source_split_0",
+ restored.getAssignedSplits().iterator().next().splitId());
+ restored.getAssignedSplits().clear();
+ assertFalse(
+ original.getAssignedSplits().isEmpty(),
+ "Defensive copy must isolate mutations on assignedSplits");
+ }
+
+ @Test
+ void testAddSplitsBackWithRegisteredReaderImmediatelyReassigns() throws
Exception {
+ TestingContext context = new TestingContext(1);
+ BigtableSourceSplitEnumerator enumerator =
+ new BigtableSourceSplitEnumerator(context, PARAMETERS);
+ enumerator.open();
+
+ context.registerReaderForTest(0);
+ enumerator.registerReader(0);
+ BigtableSourceSplit split = context.getLastAssignedSplit(0);
+
+ context.clearAssignments();
+ enumerator.addSplitsBack(Collections.singletonList(split), 0);
+
+ assertEquals(0, enumerator.currentUnassignedSplitSize());
+ assertEquals(1, context.getAssignedSplitCount(0));
+
+ BigtableSourceState checkpoint = enumerator.snapshotState(1L);
+ assertTrue(checkpoint.getPendingSplits().isEmpty());
+ assertTrue(checkpoint.getAssignedSplits().contains(split));
+ }
+
+ @SuppressWarnings("unchecked")
+ private static BigtableSourceReader createReaderWithMockedEmptyStream()
throws Exception {
+ BigtableClient mockClient = mock(BigtableClient.class);
+ BigtableDataClient mockDataClient = mock(BigtableDataClient.class);
+ when(mockClient.getDataClient()).thenReturn(mockDataClient);
+
+ SeaTunnelRowType rowType =
+ new SeaTunnelRowType(
+ new String[] {"rowkey"},
+ new
org.apache.seatunnel.api.table.type.SeaTunnelDataType[] {
+ BasicType.STRING_TYPE
+ });
+
+ BigtableSourceReader reader =
+ new BigtableSourceReader(
+ PARAMETERS, mock(SourceReader.Context.class), rowType,
mockClient);
+ reader.open();
+ return reader;
+ }
+
+ @SuppressWarnings("unchecked")
+ private static void mockEmptyReadStream(BigtableSourceReader reader,
Runnable duringForEach)
+ throws Exception {
+ BigtableClient mockClient = mock(BigtableClient.class);
+ BigtableDataClient mockDataClient = mock(BigtableDataClient.class);
+ when(mockClient.getDataClient()).thenReturn(mockDataClient);
+
+ ServerStream<com.google.cloud.bigtable.data.v2.models.Row> fakeStream =
+ mock(ServerStream.class);
+ Mockito.doAnswer(
+ invocation -> {
+ if (duringForEach != null) {
+ duringForEach.run();
+ }
+ return null;
+ })
+ .when(fakeStream)
+ .forEach(any());
+ when(mockDataClient.readRows(any(Query.class))).thenReturn(fakeStream);
+
+ // Re-bind mocked data client on the already-opened reader
+ java.lang.reflect.Field clientField =
+ BigtableSourceReader.class.getDeclaredField("bigtableClient");
+ clientField.setAccessible(true);
+ clientField.set(reader, mockClient);
+ }
+
+ private static Collector<SeaTunnelRow> mockCollector() {
+ Collector<SeaTunnelRow> collector = mock(Collector.class);
+ when(collector.getCheckpointLock()).thenReturn(new Object());
+ return collector;
+ }
+
+ private static class TestingContext
+ implements SourceSplitEnumerator.Context<BigtableSourceSplit> {
+
+ private final int parallelism;
+ private final Set<Integer> registeredReaders = new HashSet<>();
+ private final Map<Integer, List<BigtableSourceSplit>> assignments =
new HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ }
+
+ void registerReaderForTest(int subtaskId) {
+ registeredReaders.add(subtaskId);
+ }
+
+ void unregisterReaderForTest(int subtaskId) {
+ registeredReaders.remove(subtaskId);
+ }
+
+ void clearAssignments() {
+ assignments.clear();
+ }
+
+ int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+
+ BigtableSourceSplit getLastAssignedSplit(int subtaskId) {
+ List<BigtableSourceSplit> splits = assignments.get(subtaskId);
+ return splits.get(splits.size() - 1);
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return registeredReaders;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<BigtableSourceSplit>
splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+ }
+}
