[bugfix][zeta] Fixed multi-table job data loss and latency issues (ap…

…ache#149)
EricJoy2048 · Jul 6, 2023 · 0fa45a7 · 0fa45a7
1 parent 67a4a44
commit 0fa45a7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 48 deletions.
diff --git a/.../src/main/java/org/apache/seatunnel/engine/server/task/flow/ShuffleSinkFlowLifeCycle.java b/.../src/main/java/org/apache/seatunnel/engine/server/task/flow/ShuffleSinkFlowLifeCycle.java
@@ -35,8 +35,6 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
 
 @SuppressWarnings("MagicNumber")
 @Slf4j
@@ -113,43 +111,6 @@ public void close() throws IOException {
         }
     }
 
-    public CompletableFuture<Boolean> registryScheduleFlushTask(
-            ScheduledExecutorService scheduledExecutorService) {
-        // todo Register when the job started, Unload at the end(pause/cancel/crash) of the job
-        CompletableFuture<Boolean> completedFuture = new CompletableFuture();
-        Runnable scheduleFlushTask =
-                new Runnable() {
-                    @Override
-                    public void run() {
-                        if (!prepareClose
-                                && shuffleBufferSize > 0
-                                && System.currentTimeMillis() - lastModify
-                                        > shuffleBatchFlushInterval) {
-
-                            try {
-                                shuffleFlush();
-                            } catch (Exception e) {
-                                log.error("Execute schedule task error.", e);
-                            }
-                        }
-
-                        // submit next task
-                        if (!prepareClose) {
-                            Runnable nextScheduleFlushTask = this;
-                            scheduledExecutorService.schedule(
-                                    nextScheduleFlushTask,
-                                    shuffleBatchFlushInterval,
-                                    TimeUnit.MILLISECONDS);
-                        } else {
-                            completedFuture.complete(true);
-                        }
-                    }
-                };
-        scheduledExecutorService.schedule(
-                scheduleFlushTask, shuffleBatchFlushInterval, TimeUnit.MILLISECONDS);
-        return completedFuture;
-    }
-
     private synchronized void shuffleItem(Record<?> record) {
         String shuffleKey = shuffleStrategy.createShuffleKey(record, pipelineId, taskIndex);
         shuffleBuffer.computeIfAbsent(shuffleKey, key -> new LinkedList<>()).add(record);
@@ -160,8 +121,6 @@ private synchronized void shuffleItem(Record<?> record) {
                         && System.currentTimeMillis() - lastModify > shuffleBatchFlushInterval)) {
             shuffleFlush();
         }
-
-        lastModify = System.currentTimeMillis();
     }
 
     private synchronized void shuffleFlush() {
@@ -185,5 +144,6 @@ private synchronized void shuffleFlush() {
             shuffleQueueBatch.clear();
         }
         shuffleBufferSize = 0;
+        lastModify = System.currentTimeMillis();
     }
 }
diff --git a/...rc/main/java/org/apache/seatunnel/engine/server/task/flow/ShuffleSourceFlowLifeCycle.java b/...rc/main/java/org/apache/seatunnel/engine/server/task/flow/ShuffleSourceFlowLifeCycle.java
@@ -43,7 +43,7 @@ public class ShuffleSourceFlowLifeCycle<T> extends AbstractFlowLifeCycle
     private final ShuffleAction shuffleAction;
     private final int shuffleBatchSize;
     private final IQueue<Record<?>>[] shuffles;
-    private List<Record<?>> unsentBuffer;
+    private Map<Integer, List<Record<?>>> unsentBufferMap = new HashMap<>();
     private final Map<Integer, Barrier> alignedBarriers = new HashMap<>();
     private long currentCheckpointId = Long.MAX_VALUE;
     private int alignedBarriersCounter = 0;
@@ -71,6 +71,8 @@ public void collect(Collector<Record<?>> collector) throws Exception {
 
         for (int i = 0; i < shuffles.length; i++) {
             IQueue<Record<?>> shuffleQueue = shuffles[i];
+            List<Record<?>> unsentBuffer =
+                    unsentBufferMap.computeIfAbsent(i, k -> new LinkedList<>());
             if (shuffleQueue.size() == 0) {
                 emptyShuffleQueueCount++;
                 continue;
@@ -84,9 +86,9 @@ public void collect(Collector<Record<?>> collector) throws Exception {
             List<Record<?>> shuffleBatch = new LinkedList<>();
             if (alignedBarriersCounter > 0) {
                 shuffleBatch.add(shuffleQueue.take());
-            } else if (unsentBuffer != null && !unsentBuffer.isEmpty()) {
-                shuffleBatch = unsentBuffer;
-                unsentBuffer = null;
+            } else if (!unsentBuffer.isEmpty()) {
+                shuffleBatch.addAll(unsentBuffer);
+                unsentBuffer.clear();
             }
 
             shuffleQueue.drainTo(shuffleBatch, shuffleBatchSize);
@@ -121,9 +123,8 @@ public void collect(Collector<Record<?>> collector) throws Exception {
                     }
 
                     if (recordIndex + 1 < shuffleBatch.size()) {
-                        unsentBuffer =
-                                new LinkedList<>(
-                                        shuffleBatch.subList(recordIndex + 1, shuffleBatch.size()));
+                        unsentBuffer.addAll(
+                                shuffleBatch.subList(recordIndex + 1, shuffleBatch.size()));
                     }
                     break;
                 } else {