[Improve] Improve doris sink to random use be (apache#6132)

Author: Hisoka-X

seatunnel-connectors-v2/connector-doris/src/main/java/org/apache/seatunnel/connectors/doris/sink/writer/DorisSinkWriter.java

@@ -35,9 +35,7 @@
 import org.apache.seatunnel.connectors.doris.util.HttpUtil;
 import org.apache.seatunnel.connectors.doris.util.UnsupportedTypeConverterUtils;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 
 import java.io.IOException;
@@ -76,7 +74,6 @@ public class DorisSinkWriter
     private transient Thread executorThread;
     private transient volatile Exception loadException = null;
     private List<BackendV2.BackendRowV2> backends;
-    private long pos;
 
     public DorisSinkWriter(
             SinkWriter.Context context,
@@ -156,7 +153,7 @@ public void write(SeaTunnelRow element) throws IOException {
     @Override
     public Optional<DorisCommitInfo> prepareCommit() throws IOException {
         RespContent respContent = flush();
-        if (!dorisConfig.getEnable2PC()) {
+        if (!dorisConfig.getEnable2PC() || respContent == null) {
             return Optional.empty();
         }
         long txnId = respContent.getTxnId();
@@ -165,12 +162,12 @@ public Optional<DorisCommitInfo> prepareCommit() throws IOException {
                 new DorisCommitInfo(dorisStreamLoad.getHostPort(), dorisStreamLoad.getDb(), txnId));
     }
 
-    @NonNull private RespContent flush() throws IOException {
+    private RespContent flush() throws IOException {
         // disable exception checker before stop load.
         loading = false;
         checkState(dorisStreamLoad != null);
         RespContent respContent = dorisStreamLoad.stopLoad();
-        if (!DORIS_SUCCESS_STATUS.contains(respContent.getStatus())) {
+        if (respContent != null && !DORIS_SUCCESS_STATUS.contains(respContent.getStatus())) {
             String errMsg =
                     String.format(
                             "stream load error: %s, see more in %s",
@@ -253,14 +250,11 @@ public void close() throws IOException {
         }
     }
 
-    @VisibleForTesting
-    public String getAvailableBackend() {
-        long tmp = pos + backends.size();
-        while (pos < tmp) {
-            BackendV2.BackendRowV2 backend = backends.get((int) (pos % backends.size()));
+    private String getAvailableBackend() {
+        Collections.shuffle(backends);
+        for (BackendV2.BackendRowV2 backend : backends) {
             String res = backend.toBackendString();
             if (tryHttpConnection(res)) {
-                pos++;
                 return res;
             }
         }
@@ -279,7 +273,6 @@ public boolean tryHttpConnection(String backend) {
             return true;
         } catch (Exception ex) {
             log.warn("Failed to connect to backend:{}", backend, ex);
-            pos++;
             return false;
         }
     }

seatunnel-connectors-v2/connector-doris/src/main/java/org/apache/seatunnel/connectors/doris/sink/writer/DorisStreamLoad.java

@@ -78,7 +78,8 @@ public class DorisStreamLoad implements Serializable {
     private Future<CloseableHttpResponse> pendingLoadFuture;
     private final CloseableHttpClient httpClient;
     private final ExecutorService executorService;
-    private boolean loadBatchFirstRecord;
+    private volatile boolean loadBatchFirstRecord;
+    private String label;
     private long recordCount = 0;
 
     public DorisStreamLoad(
@@ -191,6 +192,8 @@ public void abortPreCommit(String labelSuffix, long chkID) throws Exception {
     public void writeRecord(byte[] record) throws IOException {
         if (loadBatchFirstRecord) {
             loadBatchFirstRecord = false;
+            recordStream.startInput();
+            startStreamLoad();
         } else {
             recordStream.write(lineDelimiter);
         }
@@ -214,21 +217,29 @@ public RespContent handlePreCommitResponse(CloseableHttpResponse response) throw
     }
 
     public RespContent stopLoad() throws IOException {
-        recordStream.endInput();
-        log.info("stream load stopped.");
-        checkState(pendingLoadFuture != null);
-        try {
-            return handlePreCommitResponse(pendingLoadFuture.get());
-        } catch (Exception e) {
-            throw new DorisConnectorException(DorisConnectorErrorCode.STREAM_LOAD_FAILED, e);
+        if (pendingLoadFuture != null) {
+            log.info("stream load stopped.");
+            recordStream.endInput();
+            try {
+                return handlePreCommitResponse(pendingLoadFuture.get());
+            } catch (Exception e) {
+                throw new DorisConnectorException(DorisConnectorErrorCode.STREAM_LOAD_FAILED, e);
+            } finally {
+                pendingLoadFuture = null;
+            }
+        } else {
+            return null;
         }
     }
 
     public void startLoad(String label) throws IOException {
         loadBatchFirstRecord = true;
         recordCount = 0;
+        this.label = label;
+    }
+
+    private void startStreamLoad() {
         HttpPutBuilder putBuilder = new HttpPutBuilder();
-        recordStream.startInput();
         log.info("stream load started for {}", label);
         try {
             InputStreamEntity entity = new InputStreamEntity(recordStream);