[optimize](sink) Optimize the BE load balancing logic during concurre…

…nt imports. (apache#388) (cherry picked from commit d9f52e6)
hyperpaas · Jan 21, 2025 · 4a88412 · 4a88412
1 parent aeea124
commit 4a88412
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/flink-doris-connector/src/main/java/org/apache/doris/flink/sink/BackendUtil.java b/flink-doris-connector/src/main/java/org/apache/doris/flink/sink/BackendUtil.java
@@ -76,9 +76,15 @@ public static BackendUtil getInstance(
     }
 
     public String getAvailableBackend() {
+        return getAvailableBackend(0);
+    }
+
+    public String getAvailableBackend(int subtaskId) {
         long tmp = pos + backends.size();
         while (pos < tmp) {
-            BackendV2.BackendRowV2 backend = backends.get((int) (pos++ % backends.size()));
+            BackendV2.BackendRowV2 backend =
+                    backends.get((int) ((pos + subtaskId) % backends.size()));
+            pos++;
             String res = backend.toBackendString();
             if (tryHttpConnection(res)) {
                 return res;

diff --git a/flink-doris-connector/src/main/java/org/apache/doris/flink/sink/writer/DorisWriter.java b/flink-doris-connector/src/main/java/org/apache/doris/flink/sink/writer/DorisWriter.java
@@ -312,7 +312,7 @@ public List<DorisWriterState> snapshotState(long checkpointId) throws IOExceptio
         List<DorisWriterState> writerStates = new ArrayList<>();
         for (DorisStreamLoad dorisStreamLoad : dorisStreamLoadMap.values()) {
             // Dynamic refresh backend
-            dorisStreamLoad.setHostPort(backendUtil.getAvailableBackend());
+            dorisStreamLoad.setHostPort(backendUtil.getAvailableBackend(subtaskId));
             DorisWriterState writerState =
                     new DorisWriterState(
                             labelPrefix,
@@ -340,7 +340,7 @@ private DorisStreamLoad getStreamLoader(String tableKey) {
                 tableKey,
                 v ->
                         new DorisStreamLoad(
-                                backendUtil.getAvailableBackend(),
+                                backendUtil.getAvailableBackend(subtaskId),
                                 dorisOptions,
                                 executionOptions,
                                 labelGenerator,
@@ -373,7 +373,7 @@ private void checkAllDone(String tableIdentifier, DorisStreamLoad dorisStreamLoa
                 // use send cached data to new txn, then notify to restart the stream
                 if (executionOptions.isUseCache()) {
                     try {
-                        dorisStreamLoad.setHostPort(backendUtil.getAvailableBackend());
+                        dorisStreamLoad.setHostPort(backendUtil.getAvailableBackend(subtaskId));
                         if (executionOptions.enabled2PC()) {
                             dorisStreamLoad.abortPreCommit(labelPrefix, curCheckpointId);
                         }

diff --git a/flink-doris-connector/src/test/java/org/apache/doris/flink/sink/TestBackendUtil.java b/flink-doris-connector/src/test/java/org/apache/doris/flink/sink/TestBackendUtil.java
@@ -22,7 +22,6 @@
 import org.junit.Ignore;
 import org.junit.Test;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
@@ -45,8 +44,7 @@ public void testGetAvailableBackend() throws Exception {
 
     @Test
     public void testTryHttpConnection() {
-        BackendUtil backendUtil = new BackendUtil(new ArrayList<>());
-        boolean flag = backendUtil.tryHttpConnection("127.0.0.1:8040");
+        boolean flag = BackendUtil.tryHttpConnection("127.0.0.1:8040");
         Assert.assertFalse(flag);
     }