[lake] Avoid create/drop connection frequently in TieringSourceReader (apache#1828)

Author: luoyuxia

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSource.java

@@ -17,6 +17,8 @@
 
 package org.apache.fluss.flink.tiering.source;
 
+import org.apache.fluss.client.Connection;
+import org.apache.fluss.client.ConnectionFactory;
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.flink.tiering.source.enumerator.TieringSourceEnumerator;
 import org.apache.fluss.flink.tiering.source.split.TieringSplit;
@@ -104,8 +106,9 @@ public SimpleVersionedSerializer<TieringSplit> getSplitSerializer() {
 
     @Override
     public SourceReader<TableBucketWriteResult<WriteResult>, TieringSplit> createReader(
-            SourceReaderContext sourceReaderContext) throws Exception {
-        return new TieringSourceReader<>(sourceReaderContext, flussConf, lakeTieringFactory);
+            SourceReaderContext sourceReaderContext) {
+        Connection connection = ConnectionFactory.createConnection(flussConf);
+        return new TieringSourceReader<>(sourceReaderContext, connection, lakeTieringFactory);
     }
 
     /** This follows the operator uid hash generation logic of flink {@link StreamGraphHasherV2}. */

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSourceReader.java

@@ -18,7 +18,7 @@
 package org.apache.fluss.flink.tiering.source;
 
 import org.apache.fluss.annotation.Internal;
-import org.apache.fluss.config.Configuration;
+import org.apache.fluss.client.Connection;
 import org.apache.fluss.flink.tiering.source.split.TieringSplit;
 import org.apache.fluss.flink.tiering.source.state.TieringSplitState;
 import org.apache.fluss.lake.writer.LakeTieringFactory;
@@ -40,15 +40,18 @@ public final class TieringSourceReader<WriteResult>
                 TieringSplit,
                 TieringSplitState> {
 
+    private final Connection connection;
+
     public TieringSourceReader(
             SourceReaderContext context,
-            Configuration flussConf,
+            Connection connection,
             LakeTieringFactory<WriteResult, ?> lakeTieringFactory) {
         super(
-                () -> new TieringSplitReader<>(flussConf, lakeTieringFactory),
+                () -> new TieringSplitReader<>(connection, lakeTieringFactory),
                 new TableBucketWriteResultEmitter<>(),
                 context.getConfiguration(),
                 context);
+        this.connection = connection;
     }
 
     @Override
@@ -85,4 +88,10 @@ protected TieringSplitState initializedState(TieringSplit split) {
     protected TieringSplit toSplitType(String splitId, TieringSplitState splitState) {
         return splitState.toSourceSplit();
     }
+
+    @Override
+    public void close() throws Exception {
+        super.close();
+        connection.close();
+    }
 }

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java

@@ -18,12 +18,10 @@
 package org.apache.fluss.flink.tiering.source;
 
 import org.apache.fluss.client.Connection;
-import org.apache.fluss.client.ConnectionFactory;
 import org.apache.fluss.client.table.Table;
 import org.apache.fluss.client.table.scanner.ScanRecord;
 import org.apache.fluss.client.table.scanner.log.LogScanner;
 import org.apache.fluss.client.table.scanner.log.ScanRecords;
-import org.apache.fluss.config.Configuration;
 import org.apache.fluss.flink.source.reader.BoundedSplitReader;
 import org.apache.fluss.flink.source.reader.RecordAndPos;
 import org.apache.fluss.flink.tiering.source.split.TieringLogSplit;
@@ -97,9 +95,10 @@ public class TieringSplitReader<WriteResult>
     private final Set<TieringLogSplit> currentTableEmptyLogSplits;
 
     public TieringSplitReader(
-            Configuration flussConf, LakeTieringFactory<WriteResult, ?> lakeTieringFactory) {
+            Connection connection, LakeTieringFactory<WriteResult, ?> lakeTieringFactory) {
         this.lakeTieringFactory = lakeTieringFactory;
-        this.connection = ConnectionFactory.createConnection(flussConf);
+        // owned by TieringSourceReader
+        this.connection = connection;
         this.pendingTieringTables = new ArrayDeque<>();
         this.pendingTieringSplits = new HashMap<>();
         this.currentTableStoppingOffsets = new HashMap<>();
@@ -457,9 +456,8 @@ public void close() throws Exception {
         if (currentTable != null) {
             currentTable.close();
         }
-        if (connection != null) {
-            connection.close();
-        }
+
+        // don't need to close connection, will be closed by TieringSourceReader
     }
 
     private void subscribeLog(TieringLogSplit logSplit) {

fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/tiering/source/TieringSplitReaderTest.java

@@ -17,6 +17,8 @@
 
 package org.apache.fluss.flink.tiering.source;
 
+import org.apache.fluss.client.Connection;
+import org.apache.fluss.client.ConnectionFactory;
 import org.apache.fluss.client.table.Table;
 import org.apache.fluss.client.table.writer.AppendWriter;
 import org.apache.fluss.client.table.writer.TableWriter;
@@ -61,7 +63,11 @@ class TieringSplitReaderTest extends FlinkTestBase {
     void testTieringTable() throws Exception {
         TablePath tablePath = TablePath.of("fluss", "fluss_test_tiering_one_table");
         long tableId = createTable(tablePath, DEFAULT_PK_TABLE_DESCRIPTOR);
-        try (TieringSplitReader<TestingWriteResult> tieringSplitReader = createTieringReader()) {
+        try (Connection connection =
+                        ConnectionFactory.createConnection(
+                                FLUSS_CLUSTER_EXTENSION.getClientConfig());
+                TieringSplitReader<TestingWriteResult> tieringSplitReader =
+                        createTieringReader(connection)) {
             // test empty splits
             SplitsAddition<TieringSplit> splitsAddition =
                     new SplitsAddition<>(
@@ -155,7 +161,11 @@ void testTieringMixTables() throws Exception {
         TablePath tablePath1 = TablePath.of("fluss", "tiering_table1");
         long tableId1 = createTable(tablePath1, DEFAULT_PK_TABLE_DESCRIPTOR);
 
-        try (TieringSplitReader<TestingWriteResult> tieringSplitReader = createTieringReader()) {
+        try (Connection connection =
+                        ConnectionFactory.createConnection(
+                                FLUSS_CLUSTER_EXTENSION.getClientConfig());
+                TieringSplitReader<TestingWriteResult> tieringSplitReader =
+                        createTieringReader(connection)) {
             Map<TableBucket, List<InternalRow>> table0Rows = putRows(tableId0, tablePath0, 10);
             Map<TableBucket, List<InternalRow>> table1Rows = putRows(tableId1, tablePath1, 10);
             waitUntilSnapshot(tableId0, 0);
@@ -266,9 +276,8 @@ void testTieringMixTables() throws Exception {
         }
     }
 
-    private TieringSplitReader<TestingWriteResult> createTieringReader() {
-        return new TieringSplitReader<>(
-                FLUSS_CLUSTER_EXTENSION.getClientConfig(), new TestingLakeTieringFactory());
+    private TieringSplitReader<TestingWriteResult> createTieringReader(Connection connection) {
+        return new TieringSplitReader<>(connection, new TestingLakeTieringFactory());
     }
 
     private void verifyTieringRows(