[improve] support combine flush async (#74)

When this is enabled, enable.combine.flush=true . Currently, streamload is still synchronous, which will block the consumption of kafka data while writing. Change it to asynchronous to improve throughput.

Author: JNSimba

src/main/java/org/apache/doris/kafka/connector/DorisSinkTask.java

@@ -109,6 +109,8 @@ public void put(final Collection<SinkRecord> records) {
                         remainingRetries);
                 remainingRetries--;
                 context.timeout(options.getRetryIntervalMs());
+                // When writing asynchronously, need to restart the asynchronous thread
+                sink.init();
                 throw new RetriableException(ex);
             }
             throw ex;

src/main/java/org/apache/doris/kafka/connector/service/DorisCombinedSinkService.java

@@ -22,8 +22,8 @@
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
+import org.apache.doris.kafka.connector.writer.AsyncStreamLoadWriter;
 import org.apache.doris.kafka.connector.writer.DorisWriter;
-import org.apache.doris.kafka.connector.writer.StreamLoadWriter;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.sink.SinkRecord;
@@ -34,14 +34,24 @@
 /** Combined all partitions and write once. */
 public class DorisCombinedSinkService extends DorisDefaultSinkService {
     private static final Logger LOG = LoggerFactory.getLogger(DorisCombinedSinkService.class);
-
     private final Map<String, HashMap<Integer, Long>> topicPartitionOffset;
 
     DorisCombinedSinkService(Map<String, String> config, SinkTaskContext context) {
         super(config, context);
         this.topicPartitionOffset = new HashMap<>();
     }
 
+    @Override
+    public void init() {
+        for (DorisWriter wr : writer.values()) {
+            if (wr instanceof AsyncStreamLoadWriter) {
+                // When the stream load asynchronous thread down,
+                // it needs to be restarted when retrying
+                ((AsyncStreamLoadWriter) wr).start();
+            }
+        }
+    }
+
     /**
      * Create new task
      *
@@ -60,8 +70,9 @@ public void startTask(final String tableName, final TopicPartition topicPartitio
             // Only by topic
             int partition = -1;
             DorisWriter dorisWriter =
-                    new StreamLoadWriter(
+                    new AsyncStreamLoadWriter(
                             tableName, topic, partition, dorisOptions, conn, connectMonitor);
+
             writer.put(writerKey, dorisWriter);
             metricsJmxReporter.start();
         }
@@ -86,14 +97,6 @@ public void insert(final Collection<SinkRecord> records) {
             // Might happen a count of record based flushing，buffer
             insert(record);
         }
-
-        // check all sink writer to see if they need to be flushed
-        for (DorisWriter writer : writer.values()) {
-            // Time based flushing
-            if (writer.shouldFlush()) {
-                writer.flushBuffer();
-            }
-        }
     }
 
     @Override
@@ -122,8 +125,9 @@ public long getOffset(final TopicPartition topicPartition) {
     public void commit(Map<TopicPartition, OffsetAndMetadata> offsets) {
         // Here we force flushing the data in memory once to
         // ensure that the offsets recorded in topicPartitionOffset have been flushed to doris
+        LOG.info("trigger flush by commit, topic {}", topicPartitionOffset.keySet());
         for (DorisWriter writer : writer.values()) {
-            writer.flushBuffer();
+            writer.commitFlush();
         }
     }
 

src/main/java/org/apache/doris/kafka/connector/service/DorisDefaultSinkService.java

@@ -87,6 +87,9 @@ public class DorisDefaultSinkService implements DorisSinkService {
                         this.metricsJmxReporter);
     }
 
+    @Override
+    public void init() {}
+
     @Override
     public void startTask(TopicPartition topicPartition) {
         startTask(dorisOptions.getTopicMapTable(topicPartition.topic()), topicPartition);

src/main/java/org/apache/doris/kafka/connector/service/DorisSinkService.java

@@ -28,6 +28,9 @@
 /** Background service of data sink, responsible to create/drop table and insert/delete files */
 public interface DorisSinkService {
 
+    /** init task for writer */
+    void init();
+
     /**
      * Start the Task.
      *

src/main/java/org/apache/doris/kafka/connector/writer/AsyncStreamLoadWriter.java

@@ -0,0 +1,139 @@
+/*
+ * 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.doris.kafka.connector.writer;
+
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import org.apache.doris.kafka.connector.cfg.DorisOptions;
+import org.apache.doris.kafka.connector.connection.ConnectionProvider;
+import org.apache.doris.kafka.connector.metrics.DorisConnectMonitor;
+import org.apache.doris.kafka.connector.utils.BackendUtils;
+import org.apache.doris.kafka.connector.writer.load.AsyncDorisStreamLoad;
+import org.apache.doris.kafka.connector.writer.load.DefaultThreadFactory;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class AsyncStreamLoadWriter extends DorisWriter {
+    private static final Logger LOG = LoggerFactory.getLogger(AsyncStreamLoadWriter.class);
+    private final LabelGenerator labelGenerator;
+    private AsyncDorisStreamLoad dorisStreamLoad;
+    private final transient ScheduledExecutorService scheduledExecutorService;
+
+    public AsyncStreamLoadWriter(
+            String tableName,
+            String topic,
+            int partition,
+            DorisOptions dorisOptions,
+            ConnectionProvider connectionProvider,
+            DorisConnectMonitor connectMonitor) {
+        super(tableName, topic, partition, dorisOptions, connectionProvider, connectMonitor);
+        this.labelGenerator = new LabelGenerator(topic, partition, tableIdentifier);
+        BackendUtils backendUtils = BackendUtils.getInstance(dorisOptions, LOG);
+        this.dorisStreamLoad =
+                new AsyncDorisStreamLoad(backendUtils, dorisOptions, topic, this.tableName);
+        this.scheduledExecutorService =
+                new ScheduledThreadPoolExecutor(
+                        1, new DefaultThreadFactory("stream-load-flush-interval"));
+        // when uploading data in streaming mode, we need to regularly detect whether there are
+        // exceptions.
+        scheduledExecutorService.scheduleWithFixedDelay(
+                this::intervalFlush,
+                dorisOptions.getFlushTime(),
+                dorisOptions.getFlushTime(),
+                TimeUnit.SECONDS);
+    }
+
+    /** start async thread stream load */
+    public void start() {
+        this.dorisStreamLoad.start();
+    }
+
+    public void insert(final SinkRecord dorisRecord) {
+        checkFlushException();
+        putBuffer(dorisRecord);
+        if (buffer.getBufferSizeBytes() >= dorisOptions.getFileSize()
+                || (dorisOptions.getRecordNum() != 0
+                        && buffer.getNumOfRecords() >= dorisOptions.getRecordNum())) {
+            LOG.info(
+                    "trigger flush by buffer size or count, buffer size: {}, num of records: {}, lastoffset : {}",
+                    buffer.getBufferSizeBytes(),
+                    buffer.getNumOfRecords(),
+                    buffer.getLastOffset());
+            bufferFullFlush();
+        }
+    }
+
+    private void bufferFullFlush() {
+        doFlush(false, true);
+    }
+
+    private void intervalFlush() {
+        LOG.debug("interval flush trigger");
+        doFlush(false, false);
+    }
+
+    public void commitFlush() {
+        doFlush(true, false);
+    }
+
+    private synchronized void doFlush(boolean waitUtilDone, boolean bufferFull) {
+        if (waitUtilDone || bufferFull) {
+            flushBuffer(waitUtilDone);
+        } else if (dorisStreamLoad.hasCapacity()) {
+            flushBuffer(false);
+        }
+    }
+
+    public synchronized void flushBuffer(boolean waitUtilDone) {
+        if (!buffer.isEmpty()) {
+            RecordBuffer tmpBuff = buffer;
+
+            String label = labelGenerator.generateLabel(tmpBuff.getLastOffset());
+            dorisStreamLoad.flush(label, tmpBuff);
+            this.buffer = new RecordBuffer();
+        }
+
+        if (waitUtilDone) {
+            dorisStreamLoad.forceFlush();
+        }
+    }
+
+    private void checkFlushException() {
+        dorisStreamLoad.checkException();
+    }
+
+    @Override
+    public void commit(int partition) {
+        // Won't go here
+    }
+
+    @Override
+    public long getOffset() {
+        // Won't go here
+        return 0;
+    }
+
+    @Override
+    public void fetchOffset() {
+        // Won't go here
+    }
+}

src/main/java/org/apache/doris/kafka/connector/writer/DorisWriter.java

@@ -42,7 +42,7 @@ public abstract class DorisWriter {
     protected String dbName;
     protected final String tableIdentifier;
     protected List<String> fileNames;
-    private RecordBuffer buffer;
+    protected RecordBuffer buffer;
     protected final AtomicLong committedOffset; // loaded offset + 1
     protected final AtomicLong flushedOffset; // flushed offset
     protected final AtomicLong processedOffset; // processed offset
@@ -116,11 +116,8 @@ protected void initRecord(final SinkRecord record) {
 
     protected void insertRecord(final SinkRecord record) {
         // discard the record if the record offset is smaller or equal to server side offset
-        // when enable.combine.flush=true, No verification is required because the offsets of
-        // multiple partitions cannot be compared.
-        if (dorisOptions.isEnableCombineFlush()
-                || (record.kafkaOffset() > this.offsetPersistedInDoris.get()
-                        && record.kafkaOffset() > processedOffset.get())) {
+        if (record.kafkaOffset() > this.offsetPersistedInDoris.get()
+                && record.kafkaOffset() > processedOffset.get()) {
             SinkRecord dorisRecord = record;
             RecordBuffer tmpBuff = null;
 
@@ -175,6 +172,11 @@ public boolean shouldFlush() {
                 >= (dorisOptions.getFlushTime() * 1000);
     }
 
+    // for combine flush
+    public void commitFlush() {
+        flushBuffer();
+    }
+
     public void flushBuffer() {
         if (buffer.isEmpty()) {
             return;

src/main/java/org/apache/doris/kafka/connector/writer/RecordBuffer.java

@@ -29,6 +29,7 @@ public class RecordBuffer extends PartitionBuffer<String> {
     private static final Logger LOG = LoggerFactory.getLogger(RecordBuffer.class);
     public static final String LINE_SEPARATOR = "\n";
     private final StringJoiner buffer;
+    private String label;
 
     public RecordBuffer() {
         super();
@@ -52,4 +53,12 @@ public String getData() {
                 getLastOffset());
         return result;
     }
+
+    public String getLabel() {
+        return label;
+    }
+
+    public void setLabel(String label) {
+        this.label = label;
+    }
 }