[Support] double buffer based async write in AppendWriteFunctionWithBufferSort

HuangZhenQiu · HuangZhenQiu · commit 1e8bf4b1d92d · 2025-09-14T21:17:25.000-07:00
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/append/AppendWriteFunctionWithBufferSort.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/append/AppendWriteFunctionWithBufferSort.java
@@ -46,6 +46,11 @@
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 /**
@@ -61,7 +66,13 @@
 public class AppendWriteFunctionWithBufferSort<T> extends AppendWriteFunction<T> {
   private static final Logger LOG = LoggerFactory.getLogger(AppendWriteFunctionWithBufferSort.class);
   private final long writeBufferSize;
-  private transient BinaryInMemorySortBuffer buffer;
+  private transient BinaryInMemorySortBuffer activeBuffer;
+  private transient BinaryInMemorySortBuffer backgroundBuffer;
+  private transient ExecutorService asyncWriteExecutor;
+  private transient AtomicReference<CompletableFuture<Void>> asyncWriteTask;
+  private transient AtomicBoolean isBackgroundBufferBeingProcessed;
+  private transient GeneratedNormalizedKeyComputer keyComputer;
+  private transient GeneratedRecordComparator recordComparator;
 
   public AppendWriteFunctionWithBufferSort(Configuration config, RowType rowType) {
     super(config, rowType);
@@ -78,41 +89,55 @@ public void open(Configuration parameters) throws Exception {
     List<String> sortKeyList = Arrays.stream(sortKeys.split(",")).map(key -> key.trim()).collect(Collectors.toList());
     SortOperatorGen sortOperatorGen = new SortOperatorGen(rowType, sortKeyList.toArray(new String[0]));
     SortCodeGenerator codeGenerator = sortOperatorGen.createSortCodeGenerator();
-    GeneratedNormalizedKeyComputer keyComputer = codeGenerator.generateNormalizedKeyComputer("SortComputer");
-    GeneratedRecordComparator recordComparator = codeGenerator.generateRecordComparator("SortComparator");
+    this.keyComputer = codeGenerator.generateNormalizedKeyComputer("SortComputer");
+    this.recordComparator = codeGenerator.generateRecordComparator("SortComparator");
     MemorySegmentPool memorySegmentPool = MemorySegmentPoolFactory.createMemorySegmentPool(config);
-    this.buffer = BufferUtils.createBuffer(rowType,
+
+    this.activeBuffer = BufferUtils.createBuffer(rowType,
             memorySegmentPool,
             keyComputer.newInstance(Thread.currentThread().getContextClassLoader()),
             recordComparator.newInstance(Thread.currentThread().getContextClassLoader()));
-    LOG.info("{} is initialized successfully.", getClass().getSimpleName());
+    this.backgroundBuffer = BufferUtils.createBuffer(rowType,
+            MemorySegmentPoolFactory.createMemorySegmentPool(config),
+            keyComputer.newInstance(Thread.currentThread().getContextClassLoader()),
+            recordComparator.newInstance(Thread.currentThread().getContextClassLoader()));
+
+    this.asyncWriteExecutor = Executors.newSingleThreadExecutor(r -> {
+      Thread t = new Thread(r, "async-write-thread");
+      t.setDaemon(true);
+      return t;
+    });
+    this.asyncWriteTask = new AtomicReference<>(CompletableFuture.completedFuture(null));
+    this.isBackgroundBufferBeingProcessed = new AtomicBoolean(false);
+
+    LOG.info("{} is initialized successfully with double buffer.", getClass().getSimpleName());
   }
 
   @Override
   public void processElement(T value, Context ctx, Collector<RowData> out) throws Exception {
     RowData data = (RowData) value;
 
-    // 1.try to write data into memory pool
-    boolean success = buffer.write(data);
+    // try to write data into active buffer
+    boolean success = activeBuffer.write(data);
     if (!success) {
-      // 2. flushes the bucket if the memory pool is full
-      sortAndSend();
-      // 3. write the row again
-      success = buffer.write(data);
+      swapAndFlushAsync();
+      // write the row again to the new active buffer
+      success = activeBuffer.write(data);
       if (!success) {
         throw new HoodieException("Buffer is too small to hold a single record.");
       }
     }
 
-    if (buffer.size() >= writeBufferSize) {
-      sortAndSend();
+    if (activeBuffer.size() >= writeBufferSize) {
+      swapAndFlushAsync();
     }
   }
 
   @Override
   public void snapshotState() {
     try {
-      sortAndSend();
+      waitForAsyncWriteCompletion();
+      sortAndSend(activeBuffer);
     } catch (IOException e) {
       throw new HoodieIOException("Fail to sort and flush data in buffer during snapshot state.", e);
     }
@@ -122,20 +147,64 @@ public void snapshotState() {
   @Override
   public void endInput() {
     try {
-      sortAndSend();
+      waitForAsyncWriteCompletion();
+      sortAndSend(activeBuffer);
     } catch (IOException e) {
       throw new HoodieIOException("Fail to sort and flush data in buffer during endInput.", e);
     }
     super.endInput();
   }
 
+  /**
+   * Swaps the active and background buffers and triggers async flush of the background buffer.
+   */
+  private void swapAndFlushAsync() throws IOException {
+    waitForAsyncWriteCompletion();
+
+    // Swap buffers
+    BinaryInMemorySortBuffer temp = activeBuffer;
+    activeBuffer = backgroundBuffer;
+    backgroundBuffer = temp;
+
+    // Start async processing of the background buffer
+    if (!backgroundBuffer.isEmpty()) {
+      isBackgroundBufferBeingProcessed.set(true);
+      CompletableFuture<Void> newTask = CompletableFuture.runAsync(() -> {
+        try {
+          sortAndSend(backgroundBuffer);
+        } catch (IOException e) {
+          LOG.error("Error during async write", e);
+          throw new RuntimeException(e);
+        } finally {
+          isBackgroundBufferBeingProcessed.set(false);
+        }
+      }, asyncWriteExecutor);
+      asyncWriteTask.set(newTask);
+    }
+  }
+
+  /**
+   * Waits for any ongoing async write operation to complete.
+   */
+  private void waitForAsyncWriteCompletion() {
+    try {
+      CompletableFuture<Void> currentTask = asyncWriteTask.get();
+      if (currentTask != null) {
+        currentTask.join();
+      }
+    } catch (Exception e) {
+      LOG.error("Error waiting for async write completion", e);
+      throw new RuntimeException(e);
+    }
+  }
+
   /**
    * For append writing, the flushing can be triggered with the following conditions:
    * 1. Checkpoint trigger, in which the current remaining data in buffer are flushed and committed.
    * 2. Binary buffer is full.
    * 3. `endInput` is called for pipelines with a bounded source.
    */
-  private void sortAndSend() throws IOException {
+  private void sortAndSend(BinaryInMemorySortBuffer buffer) throws IOException {
     if (buffer.isEmpty()) {
       return;
     }
@@ -155,4 +224,16 @@ private void sortAndSend() throws IOException {
   private static void sort(BinaryInMemorySortBuffer dataBuffer) {
     new QuickSort().sort(dataBuffer);
   }
+
+  @Override
+  public void close() throws Exception {
+    try {
+      waitForAsyncWriteCompletion();
+      if (asyncWriteExecutor != null && !asyncWriteExecutor.isShutdown()) {
+        asyncWriteExecutor.shutdown();
+      }
+    } finally {
+      super.close();
+    }
+  }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/ITTestAppendWriteFunctionWithBufferSort.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/ITTestAppendWriteFunctionWithBufferSort.java
@@ -33,6 +33,7 @@
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.data.StringData;
 import org.apache.flink.table.data.TimestampData;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -48,6 +49,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
  * Test cases for {@link AppendWriteFunctionWithBufferSort}.
@@ -175,6 +177,117 @@ public void testSortedResult() throws Exception {
     assertArrayEquals(expected.toArray(), filteredResult.toArray());
   }
 
+  @Test
+  public void testMultipleCheckpoints() throws Exception {
+    List<RowData> batch1 = Arrays.asList(
+        createRowData("uuid1", "Charlie", 35, "1970-01-01 00:00:01.123", "p1"),
+        createRowData("uuid2", "Alice", 25, "1970-01-01 00:00:01.124", "p1")
+    );
+
+    List<RowData> batch2 = Arrays.asList(
+        createRowData("uuid3", "Bob", 30, "1970-01-01 00:00:01.125", "p1"),
+        createRowData("uuid4", "Diana", 28, "1970-01-01 00:00:01.126", "p1")
+    );
+
+    TestHarness testHarness = TestWriteBase.TestHarness.instance()
+        .preparePipeline(tempFile, conf);
+
+    testHarness.consume(batch1).checkpoint(1);
+    testHarness.consume(batch2).checkpoint(2);
+    testHarness.endInput();
+
+    List<GenericRecord> actualData = TestData.readAllData(new File(conf.get(FlinkOptions.PATH)), rowType, 1);
+    assertEquals(4, actualData.size());
+  }
+
+  @Test
+  public void testLargeDatasetWithMultipleFlushes() throws Exception {
+    this.conf.set(FlinkOptions.WRITE_BUFFER_SIZE, 50L);
+
+    List<RowData> inputData = new ArrayList<>();
+    for (int i = 0; i < 500; i++) {
+      inputData.add(createRowData("uuid" + i, "Name" + (i % 10), i % 100, "1970-01-01 00:00:01.123", "p" + (i % 3)));
+    }
+
+    TestWriteBase.TestHarness.instance()
+        .preparePipeline(tempFile, conf)
+        .consume(inputData)
+        .endInput();
+
+    List<GenericRecord> actualData = TestData.readAllData(new File(conf.get(FlinkOptions.PATH)), rowType, 3);
+    assertEquals(500, actualData.size());
+  }
+
+  @Test
+  public void testSortStabilityWithDuplicateKeys() throws Exception {
+    List<RowData> inputData = Arrays.asList(
+        createRowData("uuid1", "Alice", 25, "1970-01-01 00:00:01.123", "p1"),
+        createRowData("uuid2", "Alice", 25, "1970-01-01 00:00:01.124", "p1"),
+        createRowData("uuid3", "Alice", 25, "1970-01-01 00:00:01.125", "p1"),
+        createRowData("uuid4", "Bob", 30, "1970-01-01 00:00:01.126", "p1")
+    );
+
+    TestWriteBase.TestHarness.instance()
+        .preparePipeline(tempFile, conf)
+        .consume(inputData)
+        .endInput();
+
+    List<GenericRecord> actualData = TestData.readAllData(new File(conf.get(FlinkOptions.PATH)), rowType, 1);
+    assertEquals(4, actualData.size());
+
+    List<String> filteredResult = actualData.stream()
+        .map(TestData::filterOutVariablesWithoutHudiMetadata)
+        .collect(Collectors.toList());
+
+    assertTrue(filteredResult.get(0).contains("Alice"));
+    assertTrue(filteredResult.get(1).contains("Alice"));
+    assertTrue(filteredResult.get(2).contains("Alice"));
+    assertTrue(filteredResult.get(3).contains("Bob"));
+  }
+
+  @Test
+  public void testDifferentPartitions() throws Exception {
+    List<RowData> inputData = Arrays.asList(
+        createRowData("uuid1", "Alice", 25, "1970-01-01 00:00:01.123", "p1"),
+        createRowData("uuid2", "Bob", 30, "1970-01-01 00:00:01.124", "p2"),
+        createRowData("uuid3", "Charlie", 35, "1970-01-01 00:00:01.125", "p3"),
+        createRowData("uuid4", "Diana", 28, "1970-01-01 00:00:01.126", "p1")
+    );
+
+    TestWriteBase.TestHarness.instance()
+        .preparePipeline(tempFile, conf)
+        .consume(inputData)
+        .endInput();
+
+    List<GenericRecord> actualData = TestData.readAllData(new File(conf.get(FlinkOptions.PATH)), rowType, 3);
+    assertEquals(4, actualData.size());
+  }
+
+  @Test
+  public void testConcurrentWriteScenario() throws Exception {
+    this.conf.set(FlinkOptions.WRITE_BUFFER_SIZE, 20L);
+
+    List<RowData> inputData = new ArrayList<>();
+    for (int i = 0; i < 200; i++) {
+      inputData.add(createRowData("uuid" + i, "Name" + (i % 5), i % 50, "1970-01-01 00:00:01.123", "p1"));
+    }
+
+    TestHarness testHarness = TestWriteBase.TestHarness.instance()
+        .preparePipeline(tempFile, conf);
+
+    for (int i = 0; i < inputData.size(); i += 10) {
+      List<RowData> batch = inputData.subList(i, Math.min(i + 10, inputData.size()));
+      testHarness.consume(batch);
+      if (i % 50 == 0) {
+        testHarness.checkpoint(i / 50 + 1);
+      }
+    }
+    testHarness.endInput();
+
+    List<GenericRecord> actualData = TestData.readAllData(new File(conf.get(FlinkOptions.PATH)), rowType, 1);
+    assertEquals(200, actualData.size());
+  }
+
   private GenericRowData createRowData(String uuid, String name, int age, String timestamp, String partition) {
     return GenericRowData.of(StringData.fromString(uuid), StringData.fromString(name),
         age, TimestampData.fromTimestamp(Timestamp.valueOf(timestamp)), StringData.fromString(partition));
