diff --git a/docs/content.zh/docs/connectors/flink-sources/tidb-cdc.md b/docs/content.zh/docs/connectors/flink-sources/tidb-cdc.md index 458a5d4cf45..51f03e277b7 100644 --- a/docs/content.zh/docs/connectors/flink-sources/tidb-cdc.md +++ b/docs/content.zh/docs/connectors/flink-sources/tidb-cdc.md @@ -237,15 +237,14 @@ The TiDB CDC source can work in parallel reading, because there is multiple task ### DataStream Source -The TiDB CDC connector can also be a DataStream source. You can create a SourceFunction as the following shows: - -### DataStream Source +The TiDB CDC connector can also be a DataStream source. Region-based splits are computed once by the source enumerator; restore reuses the checkpointed key ranges instead of querying PD again. ```java +import org.apache.flink.api.common.eventtime.WatermarkStrategy; import org.apache.flink.api.common.typeinfo.BasicTypeInfo; import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.connector.source.Source; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; -import org.apache.flink.streaming.api.functions.source.SourceFunction; import org.apache.flink.util.Collector; import org.apache.flink.cdc.connectors.tidb.TDBSourceOptions; @@ -261,13 +260,13 @@ public class TiDBSourceExample { public static void main(String[] args) throws Exception { - SourceFunction tidbSource = + Source tidbSource = TiDBSource.builder() .database("mydb") // set captured database .tableName("products") // set captured table .tiConf( TDBSourceOptions.getTiConfiguration( - "localhost:2399", new HashMap<>())) + "localhost:2399", null, new HashMap<>())) .snapshotEventDeserializer( new TiKVSnapshotEventDeserializationSchema() { @Override @@ -302,7 +301,9 @@ public class TiDBSourceExample { // enable checkpoint env.enableCheckpointing(3000); - env.addSource(tidbSource).print().setParallelism(1); + env.fromSource(tidbSource, WatermarkStrategy.noWatermarks(), "TiDB Source") + .print() + .setParallelism(1); env.execute("Print TiDB Snapshot + Binlog"); } diff --git a/docs/content/docs/connectors/flink-sources/tidb-cdc.md b/docs/content/docs/connectors/flink-sources/tidb-cdc.md index 0e7a78b450f..e5fb965a72e 100644 --- a/docs/content/docs/connectors/flink-sources/tidb-cdc.md +++ b/docs/content/docs/connectors/flink-sources/tidb-cdc.md @@ -237,15 +237,14 @@ The TiDB CDC source can work in parallel reading, because there is multiple task ### DataStream Source -The TiDB CDC connector can also be a DataStream source. You can create a SourceFunction as the following shows: - -### DataStream Source +The TiDB CDC connector can also be a DataStream source. Region-based splits are computed once by the source enumerator; restore reuses the checkpointed key ranges instead of querying PD again. ```java +import org.apache.flink.api.common.eventtime.WatermarkStrategy; import org.apache.flink.api.common.typeinfo.BasicTypeInfo; import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.connector.source.Source; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; -import org.apache.flink.streaming.api.functions.source.SourceFunction; import org.apache.flink.util.Collector; import org.apache.flink.cdc.connectors.tidb.TDBSourceOptions; @@ -261,13 +260,13 @@ public class TiDBSourceExample { public static void main(String[] args) throws Exception { - SourceFunction tidbSource = + Source tidbSource = TiDBSource.builder() .database("mydb") // set captured database .tableName("products") // set captured table .tiConf( TDBSourceOptions.getTiConfiguration( - "localhost:2399", new HashMap<>())) + "localhost:2399", null, new HashMap<>())) .snapshotEventDeserializer( new TiKVSnapshotEventDeserializationSchema() { @Override @@ -302,7 +301,9 @@ public class TiDBSourceExample { // enable checkpoint env.enableCheckpointing(3000); - env.addSource(tidbSource).print().setParallelism(1); + env.fromSource(tidbSource, WatermarkStrategy.noWatermarks(), "TiDB Source") + .print() + .setParallelism(1); env.execute("Print TiDB Snapshot + Binlog"); } diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiDBSource.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiDBSource.java index fa74f69ba88..a2fc35cb49d 100644 --- a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiDBSource.java +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiDBSource.java @@ -17,25 +17,128 @@ package org.apache.flink.cdc.connectors.tidb; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.api.connector.source.Source; +import org.apache.flink.api.connector.source.SourceReader; +import org.apache.flink.api.connector.source.SourceReaderContext; +import org.apache.flink.api.connector.source.SplitEnumerator; +import org.apache.flink.api.connector.source.SplitEnumeratorContext; +import org.apache.flink.api.java.typeutils.ResultTypeQueryable; +import org.apache.flink.cdc.common.annotation.PublicEvolving; +import org.apache.flink.cdc.connectors.tidb.source.enumerator.TiKVEnumeratorState; +import org.apache.flink.cdc.connectors.tidb.source.enumerator.TiKVEnumeratorStateSerializer; +import org.apache.flink.cdc.connectors.tidb.source.enumerator.TiKVSourceEnumerator; +import org.apache.flink.cdc.connectors.tidb.source.reader.TiKVSourceReader; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplitSerializer; +import org.apache.flink.cdc.connectors.tidb.table.StartupMode; import org.apache.flink.cdc.connectors.tidb.table.StartupOptions; -import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction; +import org.apache.flink.core.io.SimpleVersionedSerializer; import org.tikv.common.TiConfiguration; -/** A builder to build a SourceFunction which can read snapshot and continue to read CDC events. */ -public class TiDBSource { +/** + * The TiDB CDC {@link Source} based on FLIP-27. + * + *

The enumerator splits the captured table by TiKV regions once and assigns a contiguous + * key-range to each reader. Restore reuses checkpointed splits and does not re-query PD for region + * topology. + * + * 

{@code
+ * Source source =
+ *     TiDBSource.builder()
+ *         .database("mydb")
+ *         .tableName("products")
+ *         .tiConf(tiConf)
+ *         .snapshotEventDeserializer(...)
+ *         .changeEventDeserializer(...)
+ *         .build();
+ * env.fromSource(source, WatermarkStrategy.noWatermarks(), "TiDB Source");
+ * }
+ * + * @param the output type of the source. + */ +@PublicEvolving +public class TiDBSource + implements Source, ResultTypeQueryable { + + private static final long serialVersionUID = 1L; + + private final TiKVSnapshotEventDeserializationSchema snapshotEventDeserializationSchema; + private final TiKVChangeEventDeserializationSchema changeEventDeserializationSchema; + private final TiConfiguration tiConf; + private final StartupMode startupMode; + private final String database; + private final String tableName; + + TiDBSource( + TiKVSnapshotEventDeserializationSchema snapshotEventDeserializationSchema, + TiKVChangeEventDeserializationSchema changeEventDeserializationSchema, + TiConfiguration tiConf, + StartupMode startupMode, + String database, + String tableName) { + this.snapshotEventDeserializationSchema = snapshotEventDeserializationSchema; + this.changeEventDeserializationSchema = changeEventDeserializationSchema; + this.tiConf = tiConf; + this.startupMode = startupMode; + this.database = database; + this.tableName = tableName; + } public static Builder builder() { return new Builder<>(); } + @Override + public Boundedness getBoundedness() { + return Boundedness.CONTINUOUS_UNBOUNDED; + } + + @Override + public SourceReader createReader(SourceReaderContext readerContext) { + return new TiKVSourceReader<>( + readerContext, + snapshotEventDeserializationSchema, + changeEventDeserializationSchema, + tiConf, + startupMode); + } + + @Override + public SplitEnumerator createEnumerator( + SplitEnumeratorContext enumContext) { + return new TiKVSourceEnumerator(enumContext, tiConf, database, tableName); + } + + @Override + public SplitEnumerator restoreEnumerator( + SplitEnumeratorContext enumContext, TiKVEnumeratorState checkpoint) { + return new TiKVSourceEnumerator(enumContext, tiConf, database, tableName, checkpoint); + } + + @Override + public SimpleVersionedSerializer getSplitSerializer() { + return TiKVKeyRangeSplitSerializer.INSTANCE; + } + + @Override + public SimpleVersionedSerializer getEnumeratorCheckpointSerializer() { + return TiKVEnumeratorStateSerializer.INSTANCE; + } + + @Override + public TypeInformation getProducedType() { + return snapshotEventDeserializationSchema.getProducedType(); + } + /** Builder class of {@link TiDBSource}. */ public static class Builder { private String database; private String tableName; private StartupOptions startupOptions = StartupOptions.initial(); private TiConfiguration tiConf; - private TiKVSnapshotEventDeserializationSchema snapshotEventDeserializationSchema; private TiKVChangeEventDeserializationSchema changeEventDeserializationSchema; @@ -77,9 +180,8 @@ public Builder tiConf(TiConfiguration tiConf) { return this; } - public RichParallelSourceFunction build() { - - return new TiKVRichParallelSourceFunction<>( + public TiDBSource build() { + return new TiDBSource<>( snapshotEventDeserializationSchema, changeEventDeserializationSchema, tiConf, diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiKVRichParallelSourceFunction.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiKVRichParallelSourceFunction.java index 16c130b38c8..36316653d15 100644 --- a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiKVRichParallelSourceFunction.java +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/TiKVRichParallelSourceFunction.java @@ -63,9 +63,14 @@ import java.util.concurrent.TimeUnit; /** - * The source implementation for TiKV that read snapshot events first and then read the change - * event. + * Deprecated SourceFunction implementation. Use {@link + * org.apache.flink.cdc.connectors.tidb.TiDBSource} which splits the table once on the enumerator. + * + * @deprecated Use {@link org.apache.flink.cdc.connectors.tidb.TiDBSource} with {@code + * env.fromSource(...)}. Per-subtask region discovery can assign overlapping or gapped key + * ranges when tasks start at different times. */ +@Deprecated public class TiKVRichParallelSourceFunction extends RichParallelSourceFunction implements CheckpointListener, CheckpointedFunction, ResultTypeQueryable { diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/metrics/TiDBSourceMetrics.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/metrics/TiDBSourceMetrics.java index 1f32c0f3411..8661bd3cbb6 100644 --- a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/metrics/TiDBSourceMetrics.java +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/metrics/TiDBSourceMetrics.java @@ -17,7 +17,6 @@ package org.apache.flink.cdc.connectors.tidb.metrics; -import org.apache.flink.cdc.connectors.tidb.TiKVRichParallelSourceFunction; import org.apache.flink.metrics.Gauge; import org.apache.flink.metrics.MetricGroup; @@ -25,15 +24,15 @@ import static org.apache.flink.runtime.metrics.MetricNames.CURRENT_FETCH_EVENT_TIME_LAG; import static org.apache.flink.runtime.metrics.MetricNames.SOURCE_IDLE_TIME; -/** A collection class for handling metrics in {@link TiKVRichParallelSourceFunction}. */ +/** A collection class for handling metrics in the TiDB CDC source. */ public class TiDBSourceMetrics { private final MetricGroup metricGroup; /** - * The last record processing time, which is updated after {@link - * TiKVRichParallelSourceFunction} fetches a batch of data. It's mainly used to report metrics - * sourceIdleTime for sourceIdleTime = System.currentTimeMillis() - processTime. + * The last record processing time, which is updated after the source fetches a batch of data. + * It's mainly used to report metrics sourceIdleTime for sourceIdleTime = + * System.currentTimeMillis() - processTime. */ private long processTime = 0L; diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorState.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorState.java new file mode 100644 index 00000000000..50ad03eb41c --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorState.java @@ -0,0 +1,95 @@ +/* + * 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.flink.cdc.connectors.tidb.source.enumerator; + +import org.apache.flink.cdc.common.annotation.Internal; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +/** + * Checkpoint state of {@link TiKVSourceEnumerator}. + * + *

{@code enumerated} must stay true after the first successful PD split so restore never + * re-fetches region topology. + */ +@Internal +public class TiKVEnumeratorState { + + private final List unassignedSplits; + private final boolean enumerated; + private final int parallelism; + + public TiKVEnumeratorState(List unassignedSplits, boolean enumerated) { + this(unassignedSplits, enumerated, -1); + } + + public TiKVEnumeratorState( + List unassignedSplits, boolean enumerated, int parallelism) { + this.unassignedSplits = + Collections.unmodifiableList( + new ArrayList<>(Objects.requireNonNull(unassignedSplits))); + this.enumerated = enumerated; + this.parallelism = parallelism; + } + + public List getUnassignedSplits() { + return unassignedSplits; + } + + public boolean isEnumerated() { + return enumerated; + } + + public int getParallelism() { + return parallelism; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TiKVEnumeratorState that = (TiKVEnumeratorState) o; + return enumerated == that.enumerated + && parallelism == that.parallelism + && Objects.equals(unassignedSplits, that.unassignedSplits); + } + + @Override + public int hashCode() { + return Objects.hash(unassignedSplits, enumerated, parallelism); + } + + @Override + public String toString() { + return "TiKVEnumeratorState{enumerated=" + + enumerated + + ", parallelism=" + + parallelism + + ", unassigned=" + + unassignedSplits.size() + + '}'; + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorStateSerializer.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorStateSerializer.java new file mode 100644 index 00000000000..32adb81a102 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorStateSerializer.java @@ -0,0 +1,88 @@ +/* + * 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.flink.cdc.connectors.tidb.source.enumerator; + +import org.apache.flink.cdc.common.annotation.Internal; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplitSerializer; +import org.apache.flink.core.io.SimpleVersionedSerializer; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** Serializer for {@link TiKVEnumeratorState}. */ +@Internal +public class TiKVEnumeratorStateSerializer + implements SimpleVersionedSerializer { + + public static final TiKVEnumeratorStateSerializer INSTANCE = + new TiKVEnumeratorStateSerializer(); + + private static final int VERSION = 1; + private static final ThreadLocal SERIALIZER_CACHE = + ThreadLocal.withInitial(() -> new DataOutputSerializer(64)); + + private final TiKVKeyRangeSplitSerializer splitSerializer = + TiKVKeyRangeSplitSerializer.INSTANCE; + + private TiKVEnumeratorStateSerializer() {} + + @Override + public int getVersion() { + return VERSION; + } + + @Override + public byte[] serialize(TiKVEnumeratorState state) throws IOException { + final DataOutputSerializer out = SERIALIZER_CACHE.get(); + out.writeBoolean(state.isEnumerated()); + out.writeInt(state.getParallelism()); + final List splits = state.getUnassignedSplits(); + out.writeInt(splits.size()); + for (TiKVKeyRangeSplit split : splits) { + byte[] splitBytes = splitSerializer.serialize(split); + out.writeInt(splitBytes.length); + out.write(splitBytes); + } + final byte[] result = out.getCopyOfBuffer(); + out.clear(); + return result; + } + + @Override + public TiKVEnumeratorState deserialize(int version, byte[] serialized) throws IOException { + if (version != VERSION) { + throw new IOException("Unknown TiKVEnumeratorState version: " + version); + } + final DataInputDeserializer in = new DataInputDeserializer(serialized); + final boolean enumerated = in.readBoolean(); + final int parallelism = in.readInt(); + final int size = in.readInt(); + final List splits = new ArrayList<>(size); + for (int i = 0; i < size; i++) { + int length = in.readInt(); + byte[] splitBytes = new byte[length]; + in.readFully(splitBytes); + splits.add(splitSerializer.deserialize(splitSerializer.getVersion(), splitBytes)); + } + return new TiKVEnumeratorState(splits, enumerated, parallelism); + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVSourceEnumerator.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVSourceEnumerator.java new file mode 100644 index 00000000000..a6dd7ff2cd7 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVSourceEnumerator.java @@ -0,0 +1,256 @@ +/* + * 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.flink.cdc.connectors.tidb.source.enumerator; + +import org.apache.flink.api.connector.source.SplitEnumerator; +import org.apache.flink.api.connector.source.SplitEnumeratorContext; +import org.apache.flink.cdc.common.annotation.Internal; +import org.apache.flink.cdc.common.annotation.VisibleForTesting; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; +import org.apache.flink.cdc.connectors.tidb.table.utils.TableKeyRangeUtils; +import org.apache.flink.util.FlinkRuntimeException; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.tikv.common.TiConfiguration; +import org.tikv.common.TiSession; +import org.tikv.common.meta.TiTableInfo; +import org.tikv.kvproto.Coprocessor.KeyRange; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.TreeSet; + +/** + * Enumerator that splits a TiDB table by TiKV regions once and assigns a contiguous {@link + * TiKVKeyRangeSplit} to each reader. + * + *

Restore reuses the checkpointed splits and never queries PD for a new region layout. + */ +@Internal +public class TiKVSourceEnumerator + implements SplitEnumerator { + + private static final Logger LOG = LoggerFactory.getLogger(TiKVSourceEnumerator.class); + + private final SplitEnumeratorContext context; + @Nullable private final TiConfiguration tiConf; + private final String database; + private final String tableName; + + private final List unassignedSplits; + private final TreeSet readersAwaitingSplit; + private boolean enumerated; + private int parallelism; + + public TiKVSourceEnumerator( + SplitEnumeratorContext context, + TiConfiguration tiConf, + String database, + String tableName) { + this( + context, + tiConf, + database, + tableName, + new TiKVEnumeratorState(new ArrayList<>(), false, -1)); + } + + public TiKVSourceEnumerator( + SplitEnumeratorContext context, + @Nullable TiConfiguration tiConf, + String database, + String tableName, + TiKVEnumeratorState checkpoint) { + this.context = context; + this.tiConf = tiConf; + this.database = database; + this.tableName = tableName; + this.unassignedSplits = new ArrayList<>(checkpoint.getUnassignedSplits()); + this.enumerated = checkpoint.isEnumerated(); + this.parallelism = checkpoint.getParallelism(); + this.readersAwaitingSplit = new TreeSet<>(); + } + + @VisibleForTesting + static TiKVSourceEnumerator forRestoredSplits( + SplitEnumeratorContext context, + List unassigned, + int parallelism) { + return new TiKVSourceEnumerator( + context, + null, + "db", + "table", + new TiKVEnumeratorState(unassigned, true, parallelism)); + } + + @Override + public void start() { + if (enumerated) { + if (parallelism > 0 && parallelism != context.currentParallelism()) { + throw new FlinkRuntimeException( + String.format( + "TiDB CDC does not support changing source parallelism after splits are assigned. checkpoint parallelism=%s, current=%s", + parallelism, context.currentParallelism())); + } + LOG.info( + "Restore TiDB enumerator for {}.{}, {} unassigned split(s), skip region discovery", + database, + tableName, + unassignedSplits.size()); + return; + } + discoverSplits(); + enumerated = true; + assignSplits(); + } + + private void discoverSplits() { + if (tiConf == null) { + throw new FlinkRuntimeException( + "TiConfiguration is required to split table " + database + "." + tableName); + } + this.parallelism = context.currentParallelism(); + try (TiSession session = TiSession.create(tiConf)) { + TiTableInfo tableInfo = session.getCatalog().getTable(database, tableName); + if (tableInfo == null) { + throw new FlinkRuntimeException( + String.format("Table %s.%s does not exist.", database, tableName)); + } + long tableId = tableInfo.getId(); + List ranges = + TableKeyRangeUtils.getTableKeyRangesByRegion(session, tableId, parallelism); + unassignedSplits.clear(); + for (int i = 0; i < ranges.size(); i++) { + unassignedSplits.add(TiKVKeyRangeSplit.fromKeyRange(splitId(i), ranges.get(i))); + } + LOG.info( + "Discovered {} region-based key-range split(s) for {}.{}, tableId={}, parallelism={}", + unassignedSplits.size(), + database, + tableName, + tableId, + parallelism); + } catch (FlinkRuntimeException e) { + throw e; + } catch (Exception e) { + throw new FlinkRuntimeException( + String.format( + "Failed to split table %s.%s by TiKV regions. The job will fail rather than assign overlapping or empty ranges.", + database, tableName), + e); + } + } + + @Override + public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) { + if (!context.registeredReaders().containsKey(subtaskId)) { + return; + } + readersAwaitingSplit.add(subtaskId); + assignSplits(); + } + + @Override + public void addSplitsBack(List splits, int subtaskId) { + LOG.info("Add {} split(s) back from subtask {}", splits.size(), subtaskId); + unassignedSplits.addAll(splits); + if (context.registeredReaders().containsKey(subtaskId)) { + readersAwaitingSplit.add(subtaskId); + } + assignSplits(); + } + + @Override + public void addReader(int subtaskId) { + // Wait for the reader to request a split. + } + + private void assignSplits() { + if (!enumerated) { + return; + } + final Iterator awaiting = readersAwaitingSplit.iterator(); + while (awaiting.hasNext()) { + int reader = awaiting.next(); + if (!context.registeredReaders().containsKey(reader)) { + awaiting.remove(); + continue; + } + Optional next = takeSplitForSubtask(reader); + if (next.isPresent()) { + TiKVKeyRangeSplit split = next.get(); + LOG.info("Assign {} to subtask {}", split, reader); + context.assignSplit(split, reader); + awaiting.remove(); + } else { + LOG.info("No more splits for subtask {}", reader); + context.signalNoMoreSplits(reader); + awaiting.remove(); + } + } + } + + private Optional takeSplitForSubtask(int subtaskId) { + final String expectedId = splitId(subtaskId); + Iterator iterator = unassignedSplits.iterator(); + while (iterator.hasNext()) { + TiKVKeyRangeSplit split = iterator.next(); + if (expectedId.equals(split.splitId())) { + iterator.remove(); + return Optional.of(split); + } + } + return Optional.empty(); + } + + static String splitId(int subtaskId) { + return "tidb-" + subtaskId; + } + + @Override + public TiKVEnumeratorState snapshotState(long checkpointId) { + LOG.info( + "Enumerator snapshot checkpoint {} with enumerated={}, unassigned={}", + checkpointId, + enumerated, + unassignedSplits.size()); + return new TiKVEnumeratorState(new ArrayList<>(unassignedSplits), enumerated, parallelism); + } + + @Override + public void close() throws IOException { + // TiSession is closed after discovery. + } + + @VisibleForTesting + List getUnassignedSplits() { + return unassignedSplits; + } + + @VisibleForTesting + boolean isEnumerated() { + return enumerated; + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/reader/TiKVSourceReader.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/reader/TiKVSourceReader.java new file mode 100644 index 00000000000..c4d6715bdcf --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/reader/TiKVSourceReader.java @@ -0,0 +1,406 @@ +/* + * 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.flink.cdc.connectors.tidb.source.reader; + +import org.apache.flink.api.connector.source.ReaderOutput; +import org.apache.flink.api.connector.source.SourceReader; +import org.apache.flink.api.connector.source.SourceReaderContext; +import org.apache.flink.cdc.common.annotation.Internal; +import org.apache.flink.cdc.connectors.tidb.TiKVChangeEventDeserializationSchema; +import org.apache.flink.cdc.connectors.tidb.TiKVSnapshotEventDeserializationSchema; +import org.apache.flink.cdc.connectors.tidb.metrics.TiDBSourceMetrics; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; +import org.apache.flink.cdc.connectors.tidb.table.StartupMode; +import org.apache.flink.cdc.connectors.tidb.table.utils.TableKeyRangeUtils; +import org.apache.flink.core.io.InputStatus; +import org.apache.flink.util.Collector; +import org.apache.flink.util.FlinkRuntimeException; +import org.apache.flink.util.Preconditions; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.tikv.cdc.CDCClient; +import org.tikv.common.TiConfiguration; +import org.tikv.common.TiSession; +import org.tikv.common.key.RowKey; +import org.tikv.common.meta.TiTimestamp; +import org.tikv.kvproto.Cdcpb; +import org.tikv.kvproto.Coprocessor; +import org.tikv.kvproto.Kvrpcpb; +import org.tikv.shade.com.google.protobuf.ByteString; +import org.tikv.txn.KVClient; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +/** + * Reads snapshot events then CDC events for a single {@link TiKVKeyRangeSplit}. Runtime region + * split/merge is handled by {@link CDCClient} inside the assigned key range. + */ +@Internal +public class TiKVSourceReader implements SourceReader { + + private static final Logger LOG = LoggerFactory.getLogger(TiKVSourceReader.class); + private static final int CHANGE_EVENT_BATCH = 1000; + private static final long STREAMING_VERSION_START_EPOCH = 0L; + + private final SourceReaderContext context; + private final TiKVSnapshotEventDeserializationSchema snapshotEventDeserializationSchema; + private final TiKVChangeEventDeserializationSchema changeEventDeserializationSchema; + private final TiConfiguration tiConf; + private final StartupMode startupMode; + + private TiSession session; + private TiDBSourceMetrics sourceMetrics; + private TiKVKeyRangeSplit assignedSplit; + private Coprocessor.KeyRange keyRange; + private CDCClient cdcClient; + private ReaderOutput currentOutput; + + private long resolvedTs = TiKVKeyRangeSplit.NO_RESOLVED_TS; + private TreeMap prewrites; + private TreeMap commits; + private boolean cdcStarted; + private volatile boolean running = true; + private CompletableFuture availability; + + public TiKVSourceReader( + SourceReaderContext context, + TiKVSnapshotEventDeserializationSchema snapshotEventDeserializationSchema, + TiKVChangeEventDeserializationSchema changeEventDeserializationSchema, + TiConfiguration tiConf, + StartupMode startupMode) { + this.context = context; + this.snapshotEventDeserializationSchema = snapshotEventDeserializationSchema; + this.changeEventDeserializationSchema = changeEventDeserializationSchema; + this.tiConf = tiConf; + this.startupMode = startupMode; + this.availability = new CompletableFuture<>(); + } + + @Override + public void start() { + session = TiSession.create(tiConf); + sourceMetrics = new TiDBSourceMetrics(context.metricGroup()); + sourceMetrics.registerMetrics(); + prewrites = new TreeMap<>(); + commits = new TreeMap<>(); + maybeCreateCdcClient(); + context.sendSplitRequest(); + } + + @Override + public InputStatus pollNext(ReaderOutput output) throws Exception { + this.currentOutput = output; + if (!running) { + return InputStatus.END_OF_INPUT; + } + if (assignedSplit == null) { + if (availability.isDone()) { + availability = new CompletableFuture<>(); + } + return InputStatus.NOTHING_AVAILABLE; + } + if (assignedSplit.isEmpty()) { + if (resolvedTs == TiKVKeyRangeSplit.NO_RESOLVED_TS) { + resolvedTs = STREAMING_VERSION_START_EPOCH; + } + return idle(); + } + if (!cdcStarted) { + if (startupMode == StartupMode.INITIAL && !assignedSplit.snapshotCompleted()) { + readSnapshotEvents(output); + } else if (!assignedSplit.snapshotCompleted()) { + LOG.info("Skip snapshot read for split {}", assignedSplit.splitId()); + resolvedTs = session.getTimestamp().getVersion(); + } + startCdc(); + return InputStatus.MORE_AVAILABLE; + } + boolean emitted = pollChangeEvents(output); + return emitted ? InputStatus.MORE_AVAILABLE : idle(); + } + + private InputStatus idle() { + availability = new CompletableFuture<>(); + CompletableFuture.delayedExecutor(10, TimeUnit.MILLISECONDS) + .execute(() -> availability.complete(null)); + return InputStatus.NOTHING_AVAILABLE; + } + + private void startCdc() { + if (cdcStarted || assignedSplit.isEmpty()) { + cdcStarted = true; + return; + } + LOG.info("Start CDC for split {} from resolvedTs {}", assignedSplit.splitId(), resolvedTs); + cdcClient.start(resolvedTs); + cdcStarted = true; + } + + private void readSnapshotEvents(ReaderOutput output) throws Exception { + LOG.info("Read snapshot events for split {}", assignedSplit.splitId()); + final ReaderOutputCollector collector = new ReaderOutputCollector<>(output); + try (KVClient scanClient = session.createKVClient()) { + long startTs = session.getTimestamp().getVersion(); + ByteString start = keyRange.getStart(); + while (running) { + final List segment = + scanClient.scan(start, keyRange.getEnd(), startTs); + if (segment.isEmpty()) { + resolvedTs = startTs; + break; + } + for (final Kvrpcpb.KvPair pair : segment) { + if (TableKeyRangeUtils.isRecordKey(pair.getKey().toByteArray())) { + snapshotEventDeserializationSchema.deserialize(pair, collector); + reportMetrics(0L, startTs); + } + } + start = + RowKey.toRawKey(segment.get(segment.size() - 1).getKey()) + .next() + .toByteString(); + } + } + } + + private boolean pollChangeEvents(ReaderOutput output) throws Exception { + boolean emitted = false; + for (int i = 0; i < CHANGE_EVENT_BATCH; i++) { + final Cdcpb.Event.Row row = cdcClient.get(); + if (row == null) { + break; + } + handleRow(row); + } + if (cdcClient != null) { + try { + resolvedTs = Math.max(resolvedTs, cdcClient.getMaxResolvedTs()); + } catch (Exception e) { + LOG.debug( + "resolvedTs not available yet for split {}: {}", + assignedSplit.splitId(), + e.getMessage()); + } + } + if (!commits.isEmpty()) { + emitted = flushRows(resolvedTs, output); + } + return emitted; + } + + private void handleRow(final Cdcpb.Event.Row row) { + if (!TableKeyRangeUtils.isRecordKey(row.getKey().toByteArray())) { + return; + } + switch (row.getType()) { + case COMMITTED: + prewrites.put(RowKeyWithTs.ofStart(row), row); + commits.put(RowKeyWithTs.ofCommit(row), row); + break; + case COMMIT: + commits.put(RowKeyWithTs.ofCommit(row), row); + break; + case PREWRITE: + prewrites.put(RowKeyWithTs.ofStart(row), row); + break; + case ROLLBACK: + prewrites.remove(RowKeyWithTs.ofStart(row)); + break; + default: + LOG.warn("Unsupported row type: {}", row.getType()); + } + } + + private boolean flushRows(final long timestamp, ReaderOutput output) throws Exception { + if (output == null) { + return false; + } + boolean emitted = false; + final ReaderOutputCollector collector = new ReaderOutputCollector<>(output); + while (!commits.isEmpty() && commits.firstKey().timestamp <= timestamp) { + final Cdcpb.Event.Row commitRow = commits.pollFirstEntry().getValue(); + final Cdcpb.Event.Row prewriteRow = prewrites.remove(RowKeyWithTs.ofStart(commitRow)); + if (prewriteRow == null) { + continue; + } + changeEventDeserializationSchema.deserialize(prewriteRow, collector); + reportMetrics(prewriteRow.getStartTs(), commitRow.getCommitTs()); + emitted = true; + } + return emitted; + } + + @Override + public List snapshotState(long checkpointId) { + if (assignedSplit == null) { + return Collections.emptyList(); + } + try { + if (currentOutput != null + && !commits.isEmpty() + && resolvedTs >= STREAMING_VERSION_START_EPOCH) { + flushRows(resolvedTs, currentOutput); + } + } catch (Exception e) { + throw new FlinkRuntimeException("Failed to flush CDC rows before checkpoint", e); + } + LOG.info( + "Snapshot reader checkpoint {} for split {} at resolvedTs {}", + checkpointId, + assignedSplit.splitId(), + resolvedTs); + return Collections.singletonList(assignedSplit.withResolvedTs(resolvedTs)); + } + + @Override + public CompletableFuture isAvailable() { + return availability; + } + + @Override + public void addSplits(List splits) { + Preconditions.checkState( + assignedSplit == null, + "TiDB source reader currently supports exactly one key-range split, but already has " + + assignedSplit); + Preconditions.checkArgument( + splits.size() == 1, + "TiDB source reader currently supports exactly one key-range split, got " + + splits.size()); + assignedSplit = splits.get(0); + keyRange = assignedSplit.toKeyRange(); + resolvedTs = assignedSplit.getResolvedTs(); + maybeCreateCdcClient(); + LOG.info( + "Reader subtask {} received {}, resolvedTs={}", + context.getIndexOfSubtask(), + assignedSplit, + resolvedTs); + availability.complete(null); + } + + private void maybeCreateCdcClient() { + if (cdcClient == null + && session != null + && assignedSplit != null + && !assignedSplit.isEmpty()) { + cdcClient = new CDCClient(session, keyRange); + } + } + + @Override + public void notifyNoMoreSplits() { + if (assignedSplit == null) { + LOG.warn( + "Subtask {} received no-more-splits without a key-range split", + context.getIndexOfSubtask()); + } + availability.complete(null); + } + + @Override + public void close() throws Exception { + running = false; + availability.complete(null); + if (cdcClient != null) { + cdcClient.close(); + } + if (session != null) { + session.close(); + } + } + + private void reportMetrics(long messageTs, long fetchTs) { + long now = System.currentTimeMillis(); + sourceMetrics.recordProcessTime(now); + long messageTimestamp = TiTimestamp.extractPhysical(messageTs); + long fetchTimestamp = TiTimestamp.extractPhysical(fetchTs); + if (messageTimestamp > 0L) { + if (fetchTimestamp >= messageTimestamp) { + sourceMetrics.recordFetchDelay(fetchTimestamp - messageTimestamp); + } + sourceMetrics.recordEmitDelay(now - messageTimestamp); + } + } + + private static final class ReaderOutputCollector implements Collector { + private final ReaderOutput output; + + private ReaderOutputCollector(ReaderOutput output) { + this.output = output; + } + + @Override + public void collect(T record) { + output.collect(record); + } + + @Override + public void close() {} + } + + private static final class RowKeyWithTs implements Comparable { + private final long timestamp; + private final RowKey rowKey; + + private RowKeyWithTs(final long timestamp, final RowKey rowKey) { + this.timestamp = timestamp; + this.rowKey = rowKey; + } + + @Override + public int compareTo(final RowKeyWithTs that) { + int res = Long.compare(this.timestamp, that.timestamp); + if (res == 0) { + res = Long.compare(this.rowKey.getTableId(), that.rowKey.getTableId()); + } + if (res == 0) { + res = Long.compare(this.rowKey.getHandle(), that.rowKey.getHandle()); + } + return res; + } + + @Override + public int hashCode() { + return Objects.hash(this.timestamp, this.rowKey.getTableId(), this.rowKey.getHandle()); + } + + @Override + public boolean equals(final Object thatObj) { + if (thatObj instanceof RowKeyWithTs) { + final RowKeyWithTs that = (RowKeyWithTs) thatObj; + return this.timestamp == that.timestamp && this.rowKey.equals(that.rowKey); + } + return false; + } + + static RowKeyWithTs ofStart(final Cdcpb.Event.Row row) { + return new RowKeyWithTs(row.getStartTs(), RowKey.decode(row.getKey().toByteArray())); + } + + static RowKeyWithTs ofCommit(final Cdcpb.Event.Row row) { + return new RowKeyWithTs(row.getCommitTs(), RowKey.decode(row.getKey().toByteArray())); + } + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplit.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplit.java new file mode 100644 index 00000000000..9fe8896a437 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplit.java @@ -0,0 +1,132 @@ +/* + * 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.flink.cdc.connectors.tidb.source.split; + +import org.apache.flink.api.connector.source.SourceSplit; +import org.apache.flink.cdc.common.annotation.Internal; + +import org.tikv.common.util.KeyRangeUtils; +import org.tikv.kvproto.Coprocessor.KeyRange; +import org.tikv.shade.com.google.protobuf.ByteString; + +import java.util.Arrays; +import java.util.Objects; + +/** + * A contiguous key-range split of a TiDB table. The range is the reader's ownership of the table + * key space and is not rebound to live TiKV region ids. + */ +@Internal +public class TiKVKeyRangeSplit implements SourceSplit { + + /** Sentinel resolvedTs meaning snapshot (if enabled) has not been completed. */ + public static final long NO_RESOLVED_TS = -1L; + + private final String splitId; + private final byte[] startKey; + private final byte[] endKey; + private final long resolvedTs; + + public TiKVKeyRangeSplit(String splitId, byte[] startKey, byte[] endKey, long resolvedTs) { + this.splitId = Objects.requireNonNull(splitId, "splitId"); + this.startKey = Objects.requireNonNull(startKey, "startKey"); + this.endKey = Objects.requireNonNull(endKey, "endKey"); + this.resolvedTs = resolvedTs; + } + + public static TiKVKeyRangeSplit fromKeyRange(String splitId, KeyRange keyRange) { + return fromKeyRange(splitId, keyRange, NO_RESOLVED_TS); + } + + public static TiKVKeyRangeSplit fromKeyRange( + String splitId, KeyRange keyRange, long resolvedTs) { + return new TiKVKeyRangeSplit( + splitId, + keyRange.getStart().toByteArray(), + keyRange.getEnd().toByteArray(), + resolvedTs); + } + + @Override + public String splitId() { + return splitId; + } + + public byte[] getStartKey() { + return startKey; + } + + public byte[] getEndKey() { + return endKey; + } + + public long getResolvedTs() { + return resolvedTs; + } + + public KeyRange toKeyRange() { + return KeyRangeUtils.makeCoprocRange( + ByteString.copyFrom(startKey), ByteString.copyFrom(endKey)); + } + + public boolean isEmpty() { + return Arrays.equals(startKey, endKey); + } + + public boolean snapshotCompleted() { + return resolvedTs != NO_RESOLVED_TS; + } + + public TiKVKeyRangeSplit withResolvedTs(long newResolvedTs) { + return new TiKVKeyRangeSplit(splitId, startKey, endKey, newResolvedTs); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TiKVKeyRangeSplit that = (TiKVKeyRangeSplit) o; + return resolvedTs == that.resolvedTs + && splitId.equals(that.splitId) + && Arrays.equals(startKey, that.startKey) + && Arrays.equals(endKey, that.endKey); + } + + @Override + public int hashCode() { + int result = Objects.hash(splitId, resolvedTs); + result = 31 * result + Arrays.hashCode(startKey); + result = 31 * result + Arrays.hashCode(endKey); + return result; + } + + @Override + public String toString() { + return "TiKVKeyRangeSplit{id=" + + splitId + + ", resolvedTs=" + + resolvedTs + + ", empty=" + + isEmpty() + + '}'; + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplitSerializer.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplitSerializer.java new file mode 100644 index 00000000000..e9053d3a185 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplitSerializer.java @@ -0,0 +1,80 @@ +/* + * 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.flink.cdc.connectors.tidb.source.split; + +import org.apache.flink.cdc.common.annotation.Internal; +import org.apache.flink.core.io.SimpleVersionedSerializer; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; + +import java.io.IOException; + +/** Serializer for {@link TiKVKeyRangeSplit}. */ +@Internal +public class TiKVKeyRangeSplitSerializer implements SimpleVersionedSerializer { + + public static final TiKVKeyRangeSplitSerializer INSTANCE = new TiKVKeyRangeSplitSerializer(); + + private static final int VERSION = 1; + private static final ThreadLocal SERIALIZER_CACHE = + ThreadLocal.withInitial(() -> new DataOutputSerializer(64)); + + private TiKVKeyRangeSplitSerializer() {} + + @Override + public int getVersion() { + return VERSION; + } + + @Override + public byte[] serialize(TiKVKeyRangeSplit split) throws IOException { + final DataOutputSerializer out = SERIALIZER_CACHE.get(); + out.writeUTF(split.splitId()); + writeByteArray(out, split.getStartKey()); + writeByteArray(out, split.getEndKey()); + out.writeLong(split.getResolvedTs()); + final byte[] result = out.getCopyOfBuffer(); + out.clear(); + return result; + } + + @Override + public TiKVKeyRangeSplit deserialize(int version, byte[] serialized) throws IOException { + if (version != VERSION) { + throw new IOException("Unknown TiKVKeyRangeSplit version: " + version); + } + final DataInputDeserializer in = new DataInputDeserializer(serialized); + final String splitId = in.readUTF(); + final byte[] startKey = readByteArray(in); + final byte[] endKey = readByteArray(in); + final long resolvedTs = in.readLong(); + return new TiKVKeyRangeSplit(splitId, startKey, endKey, resolvedTs); + } + + private static void writeByteArray(DataOutputSerializer out, byte[] bytes) throws IOException { + out.writeInt(bytes.length); + out.write(bytes); + } + + private static byte[] readByteArray(DataInputDeserializer in) throws IOException { + final int length = in.readInt(); + final byte[] bytes = new byte[length]; + in.readFully(bytes); + return bytes; + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/TiDBTableSource.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/TiDBTableSource.java index f9310462548..c8f8422db11 100644 --- a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/TiDBTableSource.java +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/TiDBTableSource.java @@ -24,7 +24,7 @@ import org.apache.flink.table.connector.ChangelogMode; import org.apache.flink.table.connector.source.DynamicTableSource; import org.apache.flink.table.connector.source.ScanTableSource; -import org.apache.flink.table.connector.source.SourceFunctionProvider; +import org.apache.flink.table.connector.source.SourceProvider; import org.apache.flink.table.connector.source.abilities.SupportsReadingMetadata; import org.apache.flink.table.data.RowData; import org.apache.flink.table.types.DataType; @@ -131,7 +131,7 @@ public ScanRuntimeProvider getScanRuntimeProvider(ScanContext scanContext) { .tiConf(tiConf) .snapshotEventDeserializer(snapshotEventDeserializationSchema) .changeEventDeserializer(changeEventDeserializationSchema); - return SourceFunctionProvider.of(builder.build(), false); + return SourceProvider.of(builder.build()); } @Override diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtils.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtils.java index a76b787b985..5e4c2704266 100644 --- a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtils.java +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/main/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtils.java @@ -21,22 +21,216 @@ import org.apache.flink.shaded.guava31.com.google.common.collect.ImmutableList; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.tikv.common.TiSession; +import org.tikv.common.key.Key; import org.tikv.common.key.RowKey; import org.tikv.common.util.KeyRangeUtils; +import org.tikv.common.util.RangeSplitter; +import org.tikv.common.util.RangeSplitter.RegionTask; import org.tikv.kvproto.Coprocessor.KeyRange; +import org.tikv.shade.com.google.protobuf.ByteString; import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; import java.util.List; -/** Utils to obtain the keyRange of table. */ +/** + * Utils to obtain the keyRange of table. + * + *

Prefer splitting by TiKV regions (balanced by region count) over evenly splitting the whole + * {@code long} handle space, which causes severe skew for AUTO_INCREMENT primary keys. + */ public class TableKeyRangeUtils { + + private static final Logger LOG = LoggerFactory.getLogger(TableKeyRangeUtils.class); + + private static final Comparator START_KEY_COMPARATOR = + Comparator.comparing(range -> Key.toRawKey(range.getStart())); + + private TableKeyRangeUtils() {} + public static KeyRange getTableKeyRange(final long tableId) { return KeyRangeUtils.makeCoprocRange( RowKey.createMin(tableId).toByteString(), RowKey.createBeyondMax(tableId).toByteString()); } + /** + * Fetch regions covering the table, then partition them into {@code num} contiguous chunks + * (balanced by region count). Each chunk is merged into one {@link KeyRange}. + * + *

Must be called from a single coordinator so every subtask observes the same region + * snapshot. + */ + public static List getTableKeyRangesByRegion( + final TiSession session, final long tableId, final int num) { + Preconditions.checkArgument(num > 0, "Illegal value of num"); + Preconditions.checkNotNull(session, "TiSession must not be null"); + + final KeyRange tableRange = getTableKeyRange(tableId); + if (num == 1) { + return ImmutableList.of(tableRange); + } + + final List regionTasks = + RangeSplitter.newSplitter(session.getRegionManager()) + .splitRangeByRegion(Collections.singletonList(tableRange)); + + final List regionRanges = new ArrayList<>(); + for (RegionTask regionTask : regionTasks) { + List ranges = regionTask.getRanges(); + if (ranges == null || ranges.isEmpty()) { + continue; + } + regionRanges.addAll(ranges); + } + + LOG.info( + "TableId={} covers {} region range(s), splitting into {} parallel subtask(s)", + tableId, + regionRanges.size(), + num); + + return assignRegionRanges(regionRanges, tableRange, num); + } + + /** + * Assign region key-ranges to {@code num} subtasks by contiguous chunks and merge each chunk + * into a single key range. + * + *

Input ranges are clipped to {@code tableRange} and sorted by start key. This is required + * because {@code RangeSplitter.splitRangeByRegion} groups tasks in a HashMap and does not + * return key order. Merging unsorted ranges would produce overlapping scans across subtasks. + */ + public static List assignRegionRanges( + final List regionRanges, final KeyRange tableRange, final int num) { + Preconditions.checkArgument(num > 0, "Illegal value of num"); + Preconditions.checkNotNull(regionRanges, "regionRanges must not be null"); + Preconditions.checkNotNull(tableRange, "tableRange must not be null"); + + if (num == 1) { + return ImmutableList.of(tableRange); + } + + final List sorted = normalizeRegionRanges(regionRanges, tableRange); + if (sorted.isEmpty()) { + List empty = new ArrayList<>(num); + empty.add(tableRange); + for (int i = 1; i < num; i++) { + empty.add(emptyKeyRange(tableRange)); + } + return empty; + } + + final ImmutableList.Builder builder = ImmutableList.builder(); + final int total = sorted.size(); + + // When there are fewer regions than subtasks, give one region to each of the first + // `total` subtasks and leave the rest empty. Integer-chunking would otherwise create + // leading empty slots (e.g. 2 regions / 4 tasks -> [empty, r0, empty, r1]). + if (total <= num) { + for (int i = 0; i < num; i++) { + if (i < total) { + builder.add(sorted.get(i)); + LOG.debug("Subtask {}/{} gets region range index [{}, {})", i, num, i, i + 1); + } else { + builder.add(emptyKeyRange(tableRange)); + LOG.debug("Subtask {}/{} gets empty keyRange (no region)", i, num); + } + } + return builder.build(); + } + + for (int i = 0; i < num; i++) { + final int startIdx = (int) ((long) total * i / num); + final int endIdx = (int) ((long) total * (i + 1) / num); + List slice = new ArrayList<>(sorted.subList(startIdx, endIdx)); + List merged = KeyRangeUtils.mergeSortedRanges(slice); + builder.add(spanKeyRanges(merged)); + LOG.debug( + "Subtask {}/{} gets region ranges [{}, {}), merged into {} key range(s)", + i, + num, + startIdx, + endIdx, + merged.size()); + } + return builder.build(); + } + + /** Clip to table range, drop empties, sort by start key. Visible for testing. */ + static List normalizeRegionRanges( + final List regionRanges, final KeyRange tableRange) { + final List normalized = new ArrayList<>(regionRanges.size()); + for (KeyRange range : regionRanges) { + KeyRange clipped = intersect(range, tableRange); + if (clipped != null) { + normalized.add(clipped); + } + } + normalized.sort(START_KEY_COMPARATOR); + return normalized; + } + + /** Inclusive-start exclusive-end intersection; {@code null} if empty. */ + static KeyRange intersect(final KeyRange left, final KeyRange right) { + final Key start = maxKey(Key.toRawKey(left.getStart()), Key.toRawKey(right.getStart())); + final Key end = minKey(Key.toRawKey(left.getEnd()), Key.toRawKey(right.getEnd())); + if (start.compareTo(end) >= 0) { + return null; + } + return KeyRangeUtils.makeCoprocRange(start.toByteString(), end.toByteString()); + } + + private static Key maxKey(final Key a, final Key b) { + return a.compareTo(b) >= 0 ? a : b; + } + + private static Key minKey(final Key a, final Key b) { + return a.compareTo(b) <= 0 ? a : b; + } + + /** Create a zero-width key range that yields no scan results. */ + static KeyRange emptyKeyRange(final KeyRange tableRange) { + ByteString start = tableRange.getStart(); + return KeyRangeUtils.makeCoprocRange(start, start); + } + + /** Span a list of (usually already merged contiguous) ranges into one KeyRange. */ + static KeyRange spanKeyRanges(final List ranges) { + Preconditions.checkArgument( + ranges != null && !ranges.isEmpty(), "ranges must not be empty"); + if (ranges.size() == 1) { + return ranges.get(0); + } + return KeyRangeUtils.makeCoprocRange( + ranges.get(0).getStart(), ranges.get(ranges.size() - 1).getEnd()); + } + + /** + * @deprecated Uneven for AUTO_INCREMENT keys; use {@link #getTableKeyRangesByRegion(TiSession, + * long, int)} from a single coordinator instead. + */ + @Deprecated public static List getTableKeyRanges(final long tableId, final int num) { + return getTableKeyRangesByHandle(tableId, num); + } + + /** + * @deprecated Uneven for AUTO_INCREMENT keys; use {@link #getTableKeyRangesByRegion(TiSession, + * long, int)} from a single coordinator instead. + */ + @Deprecated + public static KeyRange getTableKeyRange(final long tableId, final int num, final int idx) { + return getTableKeyRangeByHandle(tableId, num, idx); + } + + /** Legacy handle-space split kept as fallback for callers that still need it. */ + public static List getTableKeyRangesByHandle(final long tableId, final int num) { Preconditions.checkArgument(num > 0, "Illegal value of num"); if (num == 1) { @@ -64,9 +258,10 @@ public static List getTableKeyRanges(final long tableId, final int num return builder.build(); } - public static KeyRange getTableKeyRange(final long tableId, final int num, final int idx) { + public static KeyRange getTableKeyRangeByHandle( + final long tableId, final int num, final int idx) { Preconditions.checkArgument(idx >= 0 && idx < num, "Illegal value of idx"); - return getTableKeyRanges(tableId, num).get(idx); + return getTableKeyRangesByHandle(tableId, num).get(idx); } public static boolean isRecordKey(final byte[] key) { diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorStateSerializerTest.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorStateSerializerTest.java new file mode 100644 index 00000000000..b50574a05d2 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVEnumeratorStateSerializerTest.java @@ -0,0 +1,59 @@ +/* + * 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.flink.cdc.connectors.tidb.source.enumerator; + +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; + +/** Tests for {@link TiKVEnumeratorStateSerializer}. */ +class TiKVEnumeratorStateSerializerTest { + + @Test + void testRoundTrip() throws IOException { + TiKVKeyRangeSplit split = + new TiKVKeyRangeSplit("tidb-0", new byte[] {1, 2}, new byte[] {3, 4}, 99L); + TiKVEnumeratorState state = + new TiKVEnumeratorState(Collections.singletonList(split), true, 4); + + TiKVEnumeratorStateSerializer serializer = TiKVEnumeratorStateSerializer.INSTANCE; + TiKVEnumeratorState restored = + serializer.deserialize(serializer.getVersion(), serializer.serialize(state)); + + Assertions.assertThat(restored).isEqualTo(state); + Assertions.assertThat(restored.isEnumerated()).isTrue(); + Assertions.assertThat(restored.getParallelism()).isEqualTo(4); + Assertions.assertThat(restored.getUnassignedSplits()).containsExactly(split); + } + + @Test + void testEmptyUnassigned() throws IOException { + TiKVEnumeratorState state = new TiKVEnumeratorState(Arrays.asList(), true, 2); + TiKVEnumeratorStateSerializer serializer = TiKVEnumeratorStateSerializer.INSTANCE; + TiKVEnumeratorState restored = + serializer.deserialize(serializer.getVersion(), serializer.serialize(state)); + + Assertions.assertThat(restored.getUnassignedSplits()).isEmpty(); + Assertions.assertThat(restored.isEnumerated()).isTrue(); + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVSourceEnumeratorTest.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVSourceEnumeratorTest.java new file mode 100644 index 00000000000..76f1d4d2d05 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/enumerator/TiKVSourceEnumeratorTest.java @@ -0,0 +1,210 @@ +/* + * 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.flink.cdc.connectors.tidb.source.enumerator; + +import org.apache.flink.api.connector.source.ReaderInfo; +import org.apache.flink.api.connector.source.SourceEvent; +import org.apache.flink.api.connector.source.SplitEnumeratorContext; +import org.apache.flink.api.connector.source.SplitsAssignment; +import org.apache.flink.cdc.connectors.tidb.source.split.TiKVKeyRangeSplit; +import org.apache.flink.metrics.groups.SplitEnumeratorMetricGroup; +import org.apache.flink.util.FlinkRuntimeException; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.Callable; +import java.util.function.BiConsumer; + +/** Tests for {@link TiKVSourceEnumerator}. */ +class TiKVSourceEnumeratorTest { + + @Test + void restoreSkipsRegionDiscovery() throws Exception { + RecordingEnumeratorContext context = new RecordingEnumeratorContext(2); + List unassigned = Arrays.asList(split("tidb-0"), split("tidb-1")); + TiKVSourceEnumerator enumerator = + TiKVSourceEnumerator.forRestoredSplits(context, unassigned, 2); + + enumerator.start(); + + Assertions.assertThat(enumerator.isEnumerated()).isTrue(); + Assertions.assertThat(enumerator.getUnassignedSplits()).hasSize(2); + } + + @Test + void assignsSplitMatchingSubtaskId() throws Exception { + RecordingEnumeratorContext context = new RecordingEnumeratorContext(2); + context.registerReader(0); + context.registerReader(1); + List unassigned = + new ArrayList<>(Arrays.asList(split("tidb-0"), split("tidb-1"))); + TiKVSourceEnumerator enumerator = + TiKVSourceEnumerator.forRestoredSplits(context, unassigned, 2); + enumerator.start(); + + enumerator.handleSplitRequest(1, "host"); + enumerator.handleSplitRequest(0, "host"); + + Assertions.assertThat(context.assignmentOf(1).splitId()).isEqualTo("tidb-1"); + Assertions.assertThat(context.assignmentOf(0).splitId()).isEqualTo("tidb-0"); + Assertions.assertThat(enumerator.getUnassignedSplits()).isEmpty(); + } + + @Test + void restoreDoesNotReassignSplitAlreadyOwnedByReader() throws Exception { + RecordingEnumeratorContext context = new RecordingEnumeratorContext(2); + context.registerReader(0); + // tidb-0 was already checkpointed on reader 0; enumerator only keeps tidb-1. + TiKVSourceEnumerator enumerator = + TiKVSourceEnumerator.forRestoredSplits( + context, new ArrayList<>(Arrays.asList(split("tidb-1"))), 2); + enumerator.start(); + + enumerator.handleSplitRequest(0, "host"); + + Assertions.assertThat(context.assignments).doesNotContainKey(0); + Assertions.assertThat(context.noMoreSplits).contains(0); + Assertions.assertThat(enumerator.getUnassignedSplits()).hasSize(1); + } + + @Test + void addSplitsBackRequeuesOriginalSplit() throws Exception { + RecordingEnumeratorContext context = new RecordingEnumeratorContext(1); + context.registerReader(0); + TiKVKeyRangeSplit split = split("tidb-0"); + TiKVSourceEnumerator enumerator = + TiKVSourceEnumerator.forRestoredSplits( + context, new ArrayList<>(Arrays.asList(split)), 1); + enumerator.start(); + enumerator.handleSplitRequest(0, "host"); + Assertions.assertThat(enumerator.getUnassignedSplits()).isEmpty(); + + enumerator.addSplitsBack(Arrays.asList(split.withResolvedTs(128L)), 0); + enumerator.handleSplitRequest(0, "host"); + + Assertions.assertThat(context.assignmentOf(0).getResolvedTs()).isEqualTo(128L); + } + + @Test + void snapshotStateKeepsEnumeratedFlag() throws Exception { + RecordingEnumeratorContext context = new RecordingEnumeratorContext(1); + TiKVSourceEnumerator enumerator = + TiKVSourceEnumerator.forRestoredSplits( + context, new ArrayList<>(Arrays.asList(split("tidb-0"))), 1); + enumerator.start(); + + TiKVEnumeratorState state = enumerator.snapshotState(1L); + Assertions.assertThat(state.isEnumerated()).isTrue(); + Assertions.assertThat(state.getParallelism()).isEqualTo(1); + Assertions.assertThat(state.getUnassignedSplits()).hasSize(1); + } + + @Test + void rejectsParallelismChangeOnRestore() { + RecordingEnumeratorContext context = new RecordingEnumeratorContext(4); + TiKVSourceEnumerator enumerator = + TiKVSourceEnumerator.forRestoredSplits(context, new ArrayList<>(), 2); + + Assertions.assertThatThrownBy(enumerator::start) + .isInstanceOf(FlinkRuntimeException.class) + .hasMessageContaining("does not support changing source parallelism"); + } + + private static TiKVKeyRangeSplit split(String splitId) { + return new TiKVKeyRangeSplit(splitId, new byte[] {1}, new byte[] {2}, -1L); + } + + private static final class RecordingEnumeratorContext + implements SplitEnumeratorContext { + + private final int parallelism; + private final Map readers = new HashMap<>(); + private final Map assignments = new HashMap<>(); + private final Set noMoreSplits = new HashSet<>(); + + private RecordingEnumeratorContext(int parallelism) { + this.parallelism = parallelism; + } + + void registerReader(int subtaskId) { + readers.put(subtaskId, new ReaderInfo(subtaskId, "host-" + subtaskId)); + } + + TiKVKeyRangeSplit assignmentOf(int subtaskId) { + return assignments.get(subtaskId); + } + + @Override + public SplitEnumeratorMetricGroup metricGroup() { + return null; + } + + @Override + public void sendEventToSourceReader(int subtaskId, SourceEvent event) {} + + @Override + public int currentParallelism() { + return parallelism; + } + + @Override + public Map registeredReaders() { + return readers; + } + + @Override + public void assignSplits(SplitsAssignment newSplitAssignments) { + newSplitAssignments + .assignment() + .forEach( + (subtask, splits) -> { + if (!splits.isEmpty()) { + assignments.put(subtask, splits.get(0)); + } + }); + } + + @Override + public void signalNoMoreSplits(int subtask) { + noMoreSplits.add(subtask); + } + + @Override + public void callAsync(Callable callable, BiConsumer handler) {} + + @Override + public void callAsync( + Callable callable, + BiConsumer handler, + long initialDelay, + long period) {} + + @Override + public void runInCoordinatorThread(Runnable runnable) { + runnable.run(); + } + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplitSerializerTest.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplitSerializerTest.java new file mode 100644 index 00000000000..71594ad7a70 --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/source/split/TiKVKeyRangeSplitSerializerTest.java @@ -0,0 +1,63 @@ +/* + * 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.flink.cdc.connectors.tidb.source.split; + +import org.apache.flink.cdc.connectors.tidb.table.utils.TableKeyRangeUtils; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.tikv.common.key.RowKey; +import org.tikv.common.util.KeyRangeUtils; +import org.tikv.kvproto.Coprocessor.KeyRange; + +import java.io.IOException; + +/** Tests for {@link TiKVKeyRangeSplitSerializer}. */ +class TiKVKeyRangeSplitSerializerTest { + + @Test + void testRoundTrip() throws IOException { + KeyRange range = + KeyRangeUtils.makeCoprocRange( + RowKey.toRowKey(100L, 1L).toByteString(), + RowKey.toRowKey(100L, 200L).toByteString()); + TiKVKeyRangeSplit split = TiKVKeyRangeSplit.fromKeyRange("tidb-0", range, 42L); + + TiKVKeyRangeSplitSerializer serializer = TiKVKeyRangeSplitSerializer.INSTANCE; + byte[] bytes = serializer.serialize(split); + TiKVKeyRangeSplit restored = serializer.deserialize(serializer.getVersion(), bytes); + + Assertions.assertThat(restored).isEqualTo(split); + Assertions.assertThat(restored.toKeyRange().getStart()).isEqualTo(range.getStart()); + Assertions.assertThat(restored.toKeyRange().getEnd()).isEqualTo(range.getEnd()); + } + + @Test + void testEmptyRangeRoundTrip() throws IOException { + KeyRange tableRange = TableKeyRangeUtils.getTableKeyRange(7L); + byte[] start = tableRange.getStart().toByteArray(); + TiKVKeyRangeSplit split = new TiKVKeyRangeSplit("tidb-3", start, start, 0L); + + TiKVKeyRangeSplitSerializer serializer = TiKVKeyRangeSplitSerializer.INSTANCE; + TiKVKeyRangeSplit restored = + serializer.deserialize(serializer.getVersion(), serializer.serialize(split)); + + Assertions.assertThat(restored.isEmpty()).isTrue(); + Assertions.assertThat(restored).isEqualTo(split); + } +} diff --git a/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtilsTest.java b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtilsTest.java new file mode 100644 index 00000000000..c040167701f --- /dev/null +++ b/flink-cdc-connect/flink-cdc-source-connectors/flink-connector-tidb-cdc/src/test/java/org/apache/flink/cdc/connectors/tidb/table/utils/TableKeyRangeUtilsTest.java @@ -0,0 +1,164 @@ +/* + * 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.flink.cdc.connectors.tidb.table.utils; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.tikv.common.key.Key; +import org.tikv.common.key.RowKey; +import org.tikv.common.util.KeyRangeUtils; +import org.tikv.kvproto.Coprocessor.KeyRange; +import org.tikv.shade.com.google.protobuf.ByteString; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** Tests for region-based key range assignment. */ +class TableKeyRangeUtilsTest { + + private static final long TABLE_ID = 100L; + + @Test + void testAssignRegionRangesBalancesByRegionCount() { + KeyRange tableRange = TableKeyRangeUtils.getTableKeyRange(TABLE_ID); + List regionRanges = recordRanges(0, 8); + + List assigned = + TableKeyRangeUtils.assignRegionRanges(regionRanges, tableRange, 4); + Assertions.assertThat(assigned).hasSize(4); + + // Each subtask should get 2 regions worth of handle span. + for (int i = 0; i < 4; i++) { + KeyRange range = assigned.get(i); + long expectedStart = 1_000_000L * (i * 2); + long expectedEnd = 1_000_000L * (i * 2 + 2); + Assertions.assertThat(range.getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, expectedStart).toByteString()); + Assertions.assertThat(range.getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, expectedEnd).toByteString()); + } + assertContiguousNonOverlapping(assigned); + } + + @Test + void testAssignRegionRangesSortsUnsortedInput() { + KeyRange tableRange = TableKeyRangeUtils.getTableKeyRange(TABLE_ID); + // HashMap-like insertion order: not sorted by start key. + List unsorted = + Arrays.asList( + recordRange(6_000_000L, 7_000_000L), + recordRange(1_000_000L, 2_000_000L), + recordRange(4_000_000L, 5_000_000L), + recordRange(2_000_000L, 3_000_000L)); + + List assigned = TableKeyRangeUtils.assignRegionRanges(unsorted, tableRange, 2); + Assertions.assertThat(assigned).hasSize(2); + Assertions.assertThat(assigned.get(0).getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 1_000_000L).toByteString()); + Assertions.assertThat(assigned.get(0).getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 3_000_000L).toByteString()); + Assertions.assertThat(assigned.get(1).getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 4_000_000L).toByteString()); + Assertions.assertThat(assigned.get(1).getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 7_000_000L).toByteString()); + assertContiguousNonOverlapping(assigned); + } + + @Test + void testAssignRegionRangesClipsToTableRange() { + KeyRange tableRange = + KeyRangeUtils.makeCoprocRange( + RowKey.toRowKey(TABLE_ID, 100).toByteString(), + RowKey.toRowKey(TABLE_ID, 200).toByteString()); + List regionRanges = + Arrays.asList(recordRange(0, 150), recordRange(150, 300), recordRange(300, 400)); + + List assigned = + TableKeyRangeUtils.assignRegionRanges(regionRanges, tableRange, 2); + Assertions.assertThat(assigned).hasSize(2); + Assertions.assertThat(assigned.get(0).getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 100).toByteString()); + Assertions.assertThat(assigned.get(0).getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 150).toByteString()); + Assertions.assertThat(assigned.get(1).getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 150).toByteString()); + Assertions.assertThat(assigned.get(1).getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 200).toByteString()); + assertContiguousNonOverlapping(assigned); + } + + @Test + void testAssignRegionRangesWhenParallelismExceedsRegions() { + KeyRange tableRange = TableKeyRangeUtils.getTableKeyRange(TABLE_ID); + List regionRanges = new ArrayList<>(); + regionRanges.add(recordRange(1, 100)); + regionRanges.add(recordRange(100, 200)); + + List assigned = + TableKeyRangeUtils.assignRegionRanges(regionRanges, tableRange, 4); + Assertions.assertThat(assigned).hasSize(4); + + // First two subtasks each get one region; the rest get empty ranges. + Assertions.assertThat(assigned.get(0).getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 1).toByteString()); + Assertions.assertThat(assigned.get(0).getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 100).toByteString()); + Assertions.assertThat(assigned.get(1).getStart()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 100).toByteString()); + Assertions.assertThat(assigned.get(1).getEnd()) + .isEqualTo(RowKey.toRowKey(TABLE_ID, 200).toByteString()); + Assertions.assertThat(isEmptyRange(assigned.get(2))).isTrue(); + Assertions.assertThat(isEmptyRange(assigned.get(3))).isTrue(); + } + + @Test + void testIsRecordKey() { + ByteString record = RowKey.toRowKey(TABLE_ID, 42L).toByteString(); + Assertions.assertThat(TableKeyRangeUtils.isRecordKey(record.toByteArray())).isTrue(); + } + + private static List recordRanges(int fromInclusive, int count) { + List regionRanges = new ArrayList<>(); + for (int i = fromInclusive; i < fromInclusive + count; i++) { + regionRanges.add(recordRange(1_000_000L * i, 1_000_000L * (i + 1))); + } + return regionRanges; + } + + private static KeyRange recordRange(long startHandle, long endHandle) { + return KeyRangeUtils.makeCoprocRange( + RowKey.toRowKey(TABLE_ID, startHandle).toByteString(), + RowKey.toRowKey(TABLE_ID, endHandle).toByteString()); + } + + private static boolean isEmptyRange(KeyRange range) { + return range.getStart().equals(range.getEnd()); + } + + private static void assertContiguousNonOverlapping(List ranges) { + for (int i = 1; i < ranges.size(); i++) { + if (isEmptyRange(ranges.get(i - 1)) || isEmptyRange(ranges.get(i))) { + continue; + } + Key prevEnd = Key.toRawKey(ranges.get(i - 1).getEnd()); + Key curStart = Key.toRawKey(ranges.get(i).getStart()); + Assertions.assertThat(prevEnd.compareTo(curStart)).isLessThanOrEqualTo(0); + } + } +}