[feat-34672][kudu] kudu结果表支持batchSize

FlechazoW · FlechazoW · commit b4d57b10716d · 2021-01-14T21:47:02.000+08:00
diff --git a/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/KuduOutputFormat.java b/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/KuduOutputFormat.java
@@ -18,7 +18,9 @@
 
 package com.dtstack.flink.sql.sink.kudu;
 
+import com.dtstack.flink.sql.factory.DTThreadFactory;
 import com.dtstack.flink.sql.outputformat.AbstractDtRichOutputFormat;
+import com.dtstack.flink.sql.sink.kudu.table.KuduTableInfo;
 import com.dtstack.flink.sql.util.KrbUtils;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.tuple.Tuple2;
@@ -31,6 +33,7 @@
 import org.apache.kudu.client.KuduTable;
 import org.apache.kudu.client.Operation;
 import org.apache.kudu.client.PartialRow;
+import org.apache.kudu.client.SessionConfiguration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -40,12 +43,17 @@
 import java.sql.Timestamp;
 import java.util.Date;
 import java.util.Objects;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * @author gituser
  * @modify xiuzhu
  */
-public class KuduOutputFormat extends AbstractDtRichOutputFormat<Tuple2> {
+public class KuduOutputFormat extends AbstractDtRichOutputFormat<Tuple2<Boolean, Row>> {
 
     private static final long serialVersionUID = 1L;
 
@@ -75,6 +83,24 @@ public class KuduOutputFormat extends AbstractDtRichOutputFormat<Tuple2> {
     private String keytab;
     private String krb5conf;
 
+    /**
+     * batch size
+     */
+    private Integer batchSize;
+    private Integer batchWaitInterval;
+    /**
+     * kudu session flush mode
+     */
+    private String flushMode;
+
+    private transient AtomicInteger rowCount;
+
+    /**
+     * 定时任务
+     */
+    private transient ScheduledExecutorService scheduler;
+    private transient ScheduledFuture<?> scheduledFuture;
+
     private KuduOutputFormat() {
     }
 
@@ -91,6 +117,33 @@ public void configure(Configuration parameters) {
     public void open(int taskNumber, int numTasks) throws IOException {
         establishConnection();
         initMetric();
+        initSchedulerTask();
+        rowCount = new AtomicInteger(0);
+    }
+
+    /**
+     * init the scheduler task of {@link KuduOutputFormat#flush()}
+     */
+    private void initSchedulerTask() {
+        try {
+            if (batchWaitInterval > 0) {
+                this.scheduler = new ScheduledThreadPoolExecutor(
+                        1,
+                        new DTThreadFactory("kudu-batch-flusher")
+                );
+
+                this.scheduledFuture = this.scheduler.scheduleWithFixedDelay(
+                        () -> {
+                            synchronized (KuduOutputFormat.this) {
+                                flush();
+                            }
+                        }, batchWaitInterval, batchWaitInterval, TimeUnit.MILLISECONDS
+                );
+            }
+        } catch (Exception e) {
+            LOG.error("init schedule task failed !");
+            throw new RuntimeException(e);
+        }
     }
 
     private void establishConnection() throws IOException {
@@ -127,17 +180,46 @@ private void establishConnection() throws IOException {
         }
         LOG.info("connect kudu is succeed!");
 
-        session = client.newSession();
+        session = buildSessionWithFlushMode(flushMode, client);
+    }
+
+    /**
+     * According to the different flush mode, construct different session. Detail see {@link SessionConfiguration.FlushMode}
+     *
+     * @param flushMode  flush mode
+     * @param kuduClient kudu client
+     * @return KuduSession with flush mode
+     * @throws KuduException kudu exception when session flush
+     */
+    private KuduSession buildSessionWithFlushMode(String flushMode, KuduClient kuduClient) throws KuduException {
+        KuduSession kuduSession = kuduClient.newSession();
+        if (flushMode.equalsIgnoreCase(KuduTableInfo.KuduFlushMode.MANUAL_FLUSH.name())) {
+            kuduSession.setFlushMode(SessionConfiguration.FlushMode.MANUAL_FLUSH);
+            kuduSession.setMutationBufferSpace(
+                    Integer.parseInt(String.valueOf(Math.round(batchSize * 1.2)))
+            );
+        }
+
+        if (flushMode.equalsIgnoreCase(KuduTableInfo.KuduFlushMode.AUTO_FLUSH_SYNC.name())) {
+            LOG.warn("Parameter [batchSize] will not take effect at AUTO_FLUSH_SYNC mode.");
+            kuduSession.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_SYNC);
+        }
+
+        if (flushMode.equalsIgnoreCase(KuduTableInfo.KuduFlushMode.AUTO_FLUSH_BACKGROUND.name())) {
+            LOG.warn("Unable to determine the order of data at AUTO_FLUSH_BACKGROUND mode.");
+            kuduSession.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_BACKGROUND);
+        }
+
+        return kuduSession;
     }
 
     @Override
-    public void writeRecord(Tuple2 record) throws IOException {
-        Tuple2<Boolean, Row> tupleTrans = record;
-        Boolean retract = tupleTrans.getField(0);
+    public void writeRecord(Tuple2<Boolean, Row> record) throws IOException {
+        Boolean retract = record.getField(0);
         if (!retract) {
             return;
         }
-        Row row = tupleTrans.getField(1);
+        Row row = record.getField(1);
         if (row.getArity() != fieldNames.length) {
             if (outDirtyRecords.getCount() % DIRTY_PRINT_FREQUENCY == 0) {
                 LOG.error("record insert failed ..{}", row.toString());
@@ -151,6 +233,9 @@ public void writeRecord(Tuple2 record) throws IOException {
             if (outRecords.getCount() % ROW_PRINT_FREQUENCY == 0) {
                 LOG.info("Receive data : {}", row);
             }
+            if (rowCount.getAndIncrement() >= batchSize) {
+                flush();
+            }
             session.apply(toOperation(writeMode, row));
             outRecords.inc();
         } catch (KuduException e) {
@@ -162,6 +247,22 @@ public void writeRecord(Tuple2 record) throws IOException {
         }
     }
 
+    private synchronized void flush() {
+        try {
+            if (session.isClosed()) {
+                throw new IllegalStateException("session is closed! flush data error!");
+            }
+
+            session.flush();
+            // clear
+            rowCount.set(0);
+        } catch (KuduException kuduException) {
+            LOG.error(
+                    "flush data error!", kuduException);
+            throw new RuntimeException(kuduException);
+        }
+    }
+
     @Override
     public void close() {
         if (Objects.nonNull(session) && !session.isClosed()) {
@@ -179,6 +280,14 @@ public void close() {
                 throw new IllegalArgumentException("[closeKuduClient]:" + e.getMessage());
             }
         }
+
+        if (scheduledFuture != null) {
+            scheduledFuture.cancel(false);
+        }
+
+        if (scheduler != null) {
+            scheduler.shutdownNow();
+        }
     }
 
     private Operation toOperation(WriteMode writeMode, Row row) {
@@ -345,6 +454,21 @@ public KuduOutputFormatBuilder setEnableKrb(boolean enableKrb) {
             return this;
         }
 
+        public KuduOutputFormatBuilder setBatchSize(Integer batchSize) {
+            kuduOutputFormat.batchSize = batchSize;
+            return this;
+        }
+
+        public KuduOutputFormatBuilder setBatchWaitInterval(Integer batchWaitInterval) {
+            kuduOutputFormat.batchWaitInterval = batchWaitInterval;
+            return this;
+        }
+
+        public KuduOutputFormatBuilder setFlushMode(String flushMode) {
+            kuduOutputFormat.flushMode = flushMode;
+            return this;
+        }
+
         public KuduOutputFormat finish() {
             if (kuduOutputFormat.kuduMasters == null) {
                 throw new IllegalArgumentException("No kuduMasters supplied.");
diff --git a/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/KuduSink.java b/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/KuduSink.java
@@ -34,6 +34,9 @@ public class KuduSink implements RetractStreamTableSink<Row>, Serializable, IStr
     private String principal;
     private String keytab;
     private String krb5conf;
+    private Integer batchSize;
+    private Integer batchWaitInterval;
+    private String flushMode;
 
     @Override
     public KuduSink genStreamSink(AbstractTargetTableInfo targetTableInfo) {
@@ -51,6 +54,9 @@ public KuduSink genStreamSink(AbstractTargetTableInfo targetTableInfo) {
         this.parallelism = Objects.isNull(kuduTableInfo.getParallelism()) ?
                 parallelism : kuduTableInfo.getParallelism();
 
+        this.batchSize = kuduTableInfo.getBatchSize();
+        this.batchWaitInterval = kuduTableInfo.getBatchWaitInterval();
+        this.flushMode = kuduTableInfo.getFlushMode();
         return this;
     }
 
@@ -74,6 +80,9 @@ public DataStreamSink<Tuple2<Boolean, Row>> consumeDataStream(DataStream<Tuple2<
                 .setKeytab(this.keytab)
                 .setKrb5conf(this.krb5conf)
                 .setEnableKrb(this.enableKrb)
+                .setBatchSize(this.batchSize)
+                .setBatchWaitInterval(this.batchWaitInterval)
+                .setFlushMode(this.flushMode)
                 .finish();
         RichSinkFunction richSinkFunction = new OutputFormatSinkFunction(kuduOutputFormat);
         DataStreamSink dataStreamSink = dataStream.addSink(richSinkFunction);
diff --git a/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/table/KuduSinkParser.java b/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/table/KuduSinkParser.java
@@ -2,14 +2,15 @@
 
 import com.dtstack.flink.sql.constant.PluginParamConsts;
 import com.dtstack.flink.sql.sink.kudu.KuduOutputFormat;
-import com.dtstack.flink.sql.table.AbstractTableParser;
 import com.dtstack.flink.sql.table.AbstractTableInfo;
+import com.dtstack.flink.sql.table.AbstractTableParser;
 import com.dtstack.flink.sql.util.MathUtil;
 
 import java.math.BigDecimal;
 import java.sql.Date;
 import java.sql.Timestamp;
 import java.util.Map;
+import java.util.Objects;
 
 import static com.dtstack.flink.sql.table.AbstractTableInfo.PARALLELISM_KEY;
 
@@ -27,6 +28,16 @@ public class KuduSinkParser extends AbstractTableParser {
 
     public static final String SOCKET_READ_TIMEOUT_MS = "defaultSocketReadTimeoutMs";
 
+    public static final String BATCH_SIZE_KEY = "batchSize";
+
+    public static final Integer DEFAULT_BATCH_SIZE = 1000;
+
+    public static final String BATCH_WAIT_INTERVAL_KEY = "batchWaitInterval";
+
+    public static final Integer DEFAULT_BATCH_WAIT_INTERVAL = 60 * 1000;
+
+    public static final String SESSION_FLUSH_MODE_KEY = "flushMode";
+
     @Override
     public AbstractTableInfo getTableInfo(String tableName, String fieldsInfo, Map<String, Object> props) {
         KuduTableInfo kuduTableInfo = new KuduTableInfo();
@@ -40,15 +51,27 @@ public AbstractTableInfo getTableInfo(String tableName, String fieldsInfo, Map<S
         kuduTableInfo.setWorkerCount(MathUtil.getIntegerVal(props.get(WORKER_COUNT.toLowerCase())));
         kuduTableInfo.setDefaultOperationTimeoutMs(MathUtil.getIntegerVal(props.get(OPERATION_TIMEOUT_MS.toLowerCase())));
         kuduTableInfo.setDefaultSocketReadTimeoutMs(MathUtil.getIntegerVal(props.get(SOCKET_READ_TIMEOUT_MS.toLowerCase())));
+        kuduTableInfo.setBatchSize(MathUtil.getIntegerVal(props.getOrDefault(BATCH_SIZE_KEY, DEFAULT_BATCH_SIZE)));
+        kuduTableInfo.setBatchWaitInterval(MathUtil.getIntegerVal(props.getOrDefault(BATCH_WAIT_INTERVAL_KEY, DEFAULT_BATCH_WAIT_INTERVAL)));
+
+        if (Objects.isNull(props.get(SESSION_FLUSH_MODE_KEY))) {
+            if (kuduTableInfo.getBatchSize() > 1) {
+                kuduTableInfo.setFlushMode(KuduTableInfo.KuduFlushMode.MANUAL_FLUSH.name());
+            } else {
+                kuduTableInfo.setFlushMode(KuduTableInfo.KuduFlushMode.AUTO_FLUSH_SYNC.name());
+            }
+        } else {
+            kuduTableInfo.setFlushMode(MathUtil.getString(props.get(SESSION_FLUSH_MODE_KEY)));
+        }
 
         kuduTableInfo.setPrincipal(
-            MathUtil.getString(props.get(PluginParamConsts.PRINCIPAL))
+                MathUtil.getString(props.get(PluginParamConsts.PRINCIPAL))
         );
         kuduTableInfo.setKeytab(
-            MathUtil.getString(props.get(PluginParamConsts.KEYTAB))
+                MathUtil.getString(props.get(PluginParamConsts.KEYTAB))
         );
         kuduTableInfo.setKrb5conf(
-            MathUtil.getString(props.get(PluginParamConsts.KRB5_CONF))
+                MathUtil.getString(props.get(PluginParamConsts.KRB5_CONF))
         );
         kuduTableInfo.judgeKrbEnable();
 
diff --git a/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/table/KuduTableInfo.java b/kudu/kudu-sink/src/main/java/com/dtstack/flink/sql/sink/kudu/table/KuduTableInfo.java
@@ -28,6 +28,15 @@ public class KuduTableInfo extends AbstractTargetTableInfo implements KerberosTa
     private String keytab;
     private String krb5conf;
     boolean enableKrb;
+    /**
+     * batchSize
+     */
+    private Integer batchSize;
+    private Integer batchWaitInterval;
+    /**
+     * kudu session flush mode
+     */
+    private String flushMode;
 
     public KuduTableInfo() {
         setType(CURR_TYPE);
@@ -133,4 +142,33 @@ public void setEnableKrb(boolean enableKrb) {
         this.enableKrb = enableKrb;
     }
 
+    public Integer getBatchSize() {
+        return batchSize;
+    }
+
+    public void setBatchSize(Integer batchSize) {
+        this.batchSize = batchSize;
+    }
+
+    public Integer getBatchWaitInterval() {
+        return batchWaitInterval;
+    }
+
+    public void setBatchWaitInterval(Integer batchWaitInterval) {
+        this.batchWaitInterval = batchWaitInterval;
+    }
+
+    public String getFlushMode() {
+        return flushMode;
+    }
+
+    public void setFlushMode(String flushMode) {
+        this.flushMode = flushMode;
+    }
+
+    public enum KuduFlushMode {
+        AUTO_FLUSH_SYNC,
+        AUTO_FLUSH_BACKGROUND,
+        MANUAL_FLUSH
+    }
 }
diff --git a/localTest/pom.xml b/localTest/pom.xml
@@ -77,25 +77,25 @@
 
         <dependency>
             <groupId>com.dtstack.flink</groupId>
-            <artifactId>sql.mysql</artifactId>
+            <artifactId>sql.kudu</artifactId>
             <version>1.0-SNAPSHOT</version>
         </dependency>
 
         <dependency>
             <groupId>com.dtstack.flink</groupId>
-            <artifactId>sql.side.all.mysql</artifactId>
+            <artifactId>sql.side.all.kudu</artifactId>
             <version>1.0-SNAPSHOT</version>
         </dependency>
 
         <dependency>
             <groupId>com.dtstack.flink</groupId>
-            <artifactId>sql.sink.mysql</artifactId>
+            <artifactId>sql.sink.kudu</artifactId>
             <version>1.0-SNAPSHOT</version>
         </dependency>
 
         <dependency>
             <groupId>com.dtstack.flink</groupId>
-            <artifactId>sql.side.async.mysql</artifactId>
+            <artifactId>sql.side.async.kudu</artifactId>
             <version>1.0-SNAPSHOT</version>
         </dependency>
 
