fix(flink): improve clustering resource cleanup and coverage

hudi-common/src/main/java/org/apache/hudi/common/util/collection/CloseableConcatenatingIterator.java

@@ -25,7 +25,8 @@
  * Provides closeable iterator interface over list of iterators. Consumes all records from first iterator element
  * before moving to next iterator in the list. That is concatenating elements across multiple iterators.
  */
-public class CloseableConcatenatingIterator<T> extends ConcatenatingIterator<T> {
+public class CloseableConcatenatingIterator<T> extends ConcatenatingIterator<T>
+    implements ClosableIterator<T> {
 
   public CloseableConcatenatingIterator(List<ClosableIterator<T>> iterators) {
     super(iterators);
@@ -37,4 +38,12 @@ protected Iterator<T> advanceIterator() {
     previous.close();
     return previous;
   }
+
+  @Override
+  public void close() {
+    Iterator<T> iterator;
+    while ((iterator = super.advanceIterator()) != null) {
+      ((ClosableIterator<T>) iterator).close();
+    }
+  }
 }

hudi-common/src/test/java/org/apache/hudi/common/util/collection/TestConcatenatingIterator.java

@@ -29,6 +29,9 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class TestConcatenatingIterator {
 
@@ -62,4 +65,21 @@ public void testConcatError() {
       //
     }
   }
-}
+
+  @Test
+  public void testCloseableConcatClosesRemainingIterators() {
+    ClosableIterator<Integer> first = mock(ClosableIterator.class);
+    ClosableIterator<Integer> second = mock(ClosableIterator.class);
+    when(first.hasNext()).thenReturn(true);
+    when(first.next()).thenReturn(1);
+
+    CloseableConcatenatingIterator<Integer> iterator =
+        new CloseableConcatenatingIterator<>(Arrays.asList(first, second));
+    assertEquals(1, iterator.next());
+
+    iterator.close();
+
+    verify(first).close();
+    verify(second).close();
+  }
+}

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/BulkInsertWriterHelper.java

@@ -59,7 +59,7 @@
  * Helper class for bulk insert used by Flink.
  */
 @Slf4j
-public class BulkInsertWriterHelper {
+public class BulkInsertWriterHelper implements AutoCloseable {
 
   @Getter
   protected final String instantTime;
@@ -170,6 +170,7 @@ private HoodieRowDataCreateHandle getRowCreateHandle(String partitionPath) throw
     return handles.get(partitionPath);
   }
 
+  @Override
   public void close() throws IOException {
     if (handles.isEmpty()) {
       return;
@@ -228,4 +229,3 @@ private WriteStatus closeWriteHandle(HoodieRowDataCreateHandle rowCreateHandle)
   }
 
 }
-

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/clustering/ClusteringOperator.java

@@ -87,7 +87,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -229,50 +228,55 @@ public void endInput() {
 
   private void doClustering(String instantTime, List<ClusteringOperation> clusteringOperations) throws Exception {
     clusteringMetrics.startClustering();
-    BulkInsertWriterHelper writerHelper = new BulkInsertWriterHelper(this.conf, this.table, this.writeConfig,
+    try (BulkInsertWriterHelper writerHelper = new BulkInsertWriterHelper(this.conf, this.table, this.writeConfig,
         instantTime, this.taskID, RuntimeContextUtils.getNumberOfParallelSubtasks(getRuntimeContext()),
-        RuntimeContextUtils.getAttemptNumber(getRuntimeContext()), this.rowType, true);
-
-    Iterator<RowData> iterator;
-    if (clusteringOperations.stream().anyMatch(operation -> CollectionUtils.nonEmpty(operation.getDeltaFilePaths()))) {
-      // if there are log files, we read all records into memory for a file group and apply updates.
-      iterator = readRecordsForGroupWithLogs(clusteringOperations, instantTime);
-    } else {
-      // We want to optimize reading records for case there are no log files.
-      iterator = readRecordsForGroupBaseFiles(clusteringOperations);
-    }
-
-    if (this.sortClusteringEnabled) {
-      RowDataSerializer rowDataSerializer = new RowDataSerializer(rowType);
-      BinaryExternalSorter sorter = initSorter();
-      while (iterator.hasNext()) {
-        RowData rowData = iterator.next();
-        BinaryRowData binaryRowData = rowDataSerializer.toBinaryRow(rowData).copy();
-        sorter.write(binaryRowData);
+        RuntimeContextUtils.getAttemptNumber(getRuntimeContext()), this.rowType, true)) {
+      ClosableIterator<RowData> iterator;
+      if (clusteringOperations.stream().anyMatch(operation -> CollectionUtils.nonEmpty(operation.getDeltaFilePaths()))) {
+        // if there are log files, we read all records into memory for a file group and apply updates.
+        iterator = readRecordsForGroupWithLogs(clusteringOperations, instantTime);
+      } else {
+        // We want to optimize reading records for case there are no log files.
+        iterator = readRecordsForGroupBaseFiles(clusteringOperations);
       }
 
-      BinaryRowData row = binarySerializer.createInstance();
-      while ((row = sorter.getIterator().next(row)) != null) {
-        writerHelper.write(row);
+      try (ClosableIterator<RowData> closeableIterator = iterator) {
+        if (this.sortClusteringEnabled) {
+          RowDataSerializer rowDataSerializer = new RowDataSerializer(rowType);
+          BinaryExternalSorter sorter = initSorter();
+          try {
+            while (closeableIterator.hasNext()) {
+              RowData rowData = closeableIterator.next();
+              BinaryRowData binaryRowData = rowDataSerializer.toBinaryRow(rowData).copy();
+              sorter.write(binaryRowData);
+            }
+
+            BinaryRowData row = binarySerializer.createInstance();
+            while ((row = sorter.getIterator().next(row)) != null) {
+              writerHelper.write(row);
+            }
+          } finally {
+            sorter.close();
+          }
+        } else {
+          while (closeableIterator.hasNext()) {
+            writerHelper.write(closeableIterator.next());
+          }
+        }
       }
-      sorter.close();
-    } else {
-      while (iterator.hasNext()) {
-        writerHelper.write(iterator.next());
-      }
-    }
 
-    List<WriteStatus> writeStatuses = writerHelper.getWriteStatuses(this.taskID);
-    clusteringMetrics.endClustering();
-    collector.collect(new ClusteringCommitEvent(instantTime, getFileIds(clusteringOperations), writeStatuses, this.taskID));
-    writerHelper.close();
+      List<WriteStatus> writeStatuses = writerHelper.getWriteStatuses(this.taskID);
+      clusteringMetrics.endClustering();
+      collector.collect(new ClusteringCommitEvent(instantTime, getFileIds(clusteringOperations), writeStatuses, this.taskID));
+    }
   }
 
   /**
    * Read records from baseFiles, apply updates and convert to Iterator.
    */
   @SuppressWarnings("unchecked")
-  private Iterator<RowData> readRecordsForGroupWithLogs(List<ClusteringOperation> clusteringOps, String instantTime) {
+  private ClosableIterator<RowData> readRecordsForGroupWithLogs(
+      List<ClusteringOperation> clusteringOps, String instantTime) {
     List<ClosableIterator<RowData>> recordIterators = new ArrayList<>();
     long maxMemoryPerCompaction = MergeUtils.getMaxMemoryPerCompaction(new FlinkTaskContextSupplier(null), writeConfig);
     log.info("MaxMemoryPerCompaction run as part of clustering => {}", maxMemoryPerCompaction);
@@ -304,7 +308,7 @@ private ClosableIterator<RowData> getRecordIterator(ClusteringOperation clusterO
   /**
    * Read records from baseFiles and get iterator.
    */
-  private Iterator<RowData> readRecordsForGroupBaseFiles(List<ClusteringOperation> clusteringOps) {
+  private ClosableIterator<RowData> readRecordsForGroupBaseFiles(List<ClusteringOperation> clusteringOps) {
     List<ClosableIterator<RowData>> iteratorsForPartition = clusteringOps.stream().map(clusteringOp -> {
       try {
         HoodieFileReaderFactory fileReaderFactory = HoodieIOFactory.getIOFactory(table.getStorage())