fix(flink): close CDC image spillable maps on failures

hudi-common/src/main/java/org/apache/hudi/common/util/CloseableUtils.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.common.util;
+
+/** Utility methods for closing resources. */
+public final class CloseableUtils {
+
+  private CloseableUtils() {
+  }
+
+  /** Closes {@code closeable}, attaching any failure to {@code primary} as a suppressed exception. */
+  public static void closeSuppressing(AutoCloseable closeable, Throwable primary) {
+    try {
+      closeable.close();
+    } catch (Throwable closeError) {
+      primary.addSuppressed(closeError);
+    }
+  }
+}

hudi-common/src/test/java/org/apache/hudi/common/util/TestCloseableUtils.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.common.util;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+
+class TestCloseableUtils {
+
+  @Test
+  void testCloseSuppressing() {
+    IOException primary = new IOException("primary");
+    IOException closeError = new IOException("close");
+
+    CloseableUtils.closeSuppressing(() -> {
+      throw closeError;
+    }, primary);
+
+    assertArrayEquals(new Throwable[] {closeError}, primary.getSuppressed());
+  }
+}

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/reader/function/HoodieCdcSplitReaderFunction.java

@@ -66,6 +66,8 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.apache.hudi.common.util.CloseableUtils.closeSuppressing;
+
 /**
  * CDC reader function for source V2. Reads CDC splits ({@link HoodieCdcSourceSplit}) and
  * emits change-log {@link RowData} records tagged with the appropriate {@link org.apache.flink.types.RowKind}.
@@ -223,10 +225,15 @@ conf, getHadoopConf(), tablePath, tableSchema, requiredSchema,
         String logFilePath = new Path(tablePath, fileSplit.getCdcFiles().get(0)).toString();
         MergeOnReadInputSplit split = CdcIterators.singleLogFile2Split(tablePath, logFilePath, maxCompactionMemoryInBytes);
         ClosableIterator<HoodieRecord<RowData>> recordIterator = getFileSliceHoodieRecordIterator(split);
-        return new CdcIterators.DataLogFileIterator(
-            maxCompactionMemoryInBytes, imageManager, fileSplit, tableSchema,
-            tableState.getRequiredRowType(), tableState.getRequiredPositions(),
-            recordIterator, getMetaClient(), getWriteConfig());
+        try {
+          return new CdcIterators.DataLogFileIterator(
+              maxCompactionMemoryInBytes, imageManager, fileSplit, tableSchema,
+              tableState.getRequiredRowType(), tableState.getRequiredPositions(),
+              recordIterator, getMetaClient(), getWriteConfig());
+        } catch (IOException | RuntimeException | Error e) {
+          closeSuppressing(recordIterator, e);
+          throw e;
+        }
       }
       case REPLACE_COMMIT: {
         return new CdcIterators.ReplaceCommitIterator(

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/reader/function/HoodieSplitReaderFunction.java

@@ -44,6 +44,8 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
+import static org.apache.hudi.common.util.CloseableUtils.closeSuppressing;
+
 /**
  * Default reader function implementation for both MOR and COW tables.
  */
@@ -89,15 +91,6 @@ protected ClosableIterator<RowData> createRecordIterator(HoodieSourceSplit split
     }
   }
 
-  /** Closes {@code reader}, attaching any close failure to {@code primary} as a suppressed exception. */
-  private static void closeSuppressing(HoodieRecordReader<RowData> reader, Throwable primary) {
-    try {
-      reader.close();
-    } catch (Exception closeError) {
-      primary.addSuppressed(closeError);
-    }
-  }
-
   @Override
   protected RowType producedRowType() {
     return HoodieSchemaConverter.convertToRowType(requiredSchema);

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/format/cdc/CdcImageManager.java

@@ -45,6 +45,7 @@
 import java.util.TreeMap;
 import java.util.function.Function;
 
+import static org.apache.hudi.common.util.CloseableUtils.closeSuppressing;
 import static org.apache.hudi.hadoop.utils.HoodieInputFormatUtils.HOODIE_RECORD_KEY_COL_POS;
 
 /**
@@ -105,13 +106,16 @@ private ExternalSpillableMap<String, byte[]> loadImageRecords(
         serializer.serialize(row, new BytesArrayOutputView(baos));
         imageRecordsMap.put(recordKey, baos.toByteArray());
       }
+    } catch (IOException | RuntimeException | Error e) {
+      closeSuppressing(imageRecordsMap, e);
+      throw e;
     }
     return imageRecordsMap;
   }
 
   public RowData getImageRecord(
       String recordKey,
-      ExternalSpillableMap<String, byte[]> imageCache,
+      Map<String, byte[]> imageCache,
       RowKind rowKind) {
     byte[] bytes = imageCache.get(recordKey);
     ValidationUtils.checkState(bytes != null,
@@ -127,7 +131,7 @@ public RowData getImageRecord(
 
   public void updateImageRecord(
       String recordKey,
-      ExternalSpillableMap<String, byte[]> imageCache,
+      Map<String, byte[]> imageCache,
       RowData row) {
     ByteArrayOutputStream baos = new ByteArrayOutputStream(4096);
     try {
@@ -140,7 +144,7 @@ public void updateImageRecord(
 
   public RowData removeImageRecord(
       String recordKey,
-      ExternalSpillableMap<String, byte[]> imageCache) {
+      Map<String, byte[]> imageCache) {
     byte[] bytes = imageCache.remove(recordKey);
     if (bytes == null) {
       return null;
@@ -154,8 +158,25 @@ public RowData removeImageRecord(
 
   @Override
   public void close() {
-    cache.values().forEach(ExternalSpillableMap::close);
-    cache.clear();
+    RuntimeException failure = null;
+    try {
+      for (ExternalSpillableMap<String, byte[]> spillableMap : cache.values()) {
+        try {
+          spillableMap.close();
+        } catch (RuntimeException e) {
+          if (failure == null) {
+            failure = e;
+          } else {
+            failure.addSuppressed(e);
+          }
+        }
+      }
+    } finally {
+      cache.clear();
+    }
+    if (failure != null) {
+      throw failure;
+    }
   }
 
   // -------------------------------------------------------------------------

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/format/cdc/CdcInputFormat.java

@@ -48,6 +48,8 @@
 import java.util.List;
 import java.util.function.Function;
 
+import static org.apache.hudi.common.util.CloseableUtils.closeSuppressing;
+
 /**
  * The base InputFormat class to read Hoodie data set as change logs.
  */
@@ -157,11 +159,16 @@ private ClosableIterator<RowData> getRecordIterator(
         String logFilepath = new Path(tablePath, fileSplit.getCdcFiles().get(0)).toString();
         MergeOnReadInputSplit split = CdcIterators.singleLogFile2Split(tablePath, logFilepath, maxCompactionMemoryInBytes);
         ClosableIterator<HoodieRecord<RowData>> recordIterator = getSplitRecordIterator(split);
-        return new CdcIterators.DataLogFileIterator(
-            maxCompactionMemoryInBytes, imageManager, fileSplit,
-            HoodieSchema.parse(tableState.getTableSchema()),
-            tableState.getRequiredRowType(), tableState.getRequiredPositions(),
-            recordIterator, metaClient, imageManager.getWriteConfig());
+        try {
+          return new CdcIterators.DataLogFileIterator(
+              maxCompactionMemoryInBytes, imageManager, fileSplit,
+              HoodieSchema.parse(tableState.getTableSchema()),
+              tableState.getRequiredRowType(), tableState.getRequiredPositions(),
+              recordIterator, metaClient, imageManager.getWriteConfig());
+        } catch (IOException | RuntimeException | Error e) {
+          closeSuppressing(recordIterator, e);
+          throw e;
+        }
       case REPLACE_COMMIT:
         return new CdcIterators.ReplaceCommitIterator(
             tablePath, tableState.getRequiredRowType(), tableState.getRequiredPositions(),

hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/format/cdc/CdcIterators.java

@@ -60,7 +60,6 @@
 import org.apache.hudi.storage.HoodieStorage;
 import org.apache.hudi.storage.StoragePath;
 import org.apache.hudi.table.format.FlinkReaderContextFactory;
-import org.apache.hudi.table.format.FormatUtils;
 import org.apache.hudi.table.format.HoodieRowDataFileReader;
 import org.apache.hudi.table.format.InternalSchemaManager;
 import org.apache.hudi.table.format.mor.MergeOnReadInputSplit;
@@ -80,9 +79,11 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.apache.hudi.common.util.CloseableUtils.closeSuppressing;
 import static org.apache.hudi.table.format.FormatUtils.buildAvroRecordBySchema;
 
 /**
@@ -146,11 +147,12 @@ public RowData next() {
 
     @Override
     public void close() {
-      if (recordIterator != null) {
-        recordIterator.close();
-      }
-      if (imageManager != null) {
-        imageManager.close();
+      try (CdcImageManager ignored = imageManager) {
+        if (recordIterator != null) {
+          recordIterator.close();
+        }
+      } finally {
+        recordIterator = null;
         imageManager = null;
       }
     }
@@ -253,7 +255,7 @@ public static class DataLogFileIterator implements ClosableIterator<RowData> {
     private final String[] orderingFields;
     private final TypedProperties props;
 
-    private ExternalSpillableMap<String, byte[]> beforeImages;
+    private Map<String, byte[]> beforeImages;
     private RowData currentImage;
     private RowData sideImage;
 
@@ -287,15 +289,15 @@ public DataLogFileIterator(
           metaClient.getTableConfig().getPartialUpdateMode());
       this.logRecordIterator = logRecordIterator;
       this.deleteContext = new DeleteContext(props, tableSchema).withReaderSchema(tableSchema);
-      initImages(cdcFileSplit, writeConfig);
+      initImages(cdcFileSplit);
     }
 
-    private void initImages(HoodieCDCFileSplit fileSplit, HoodieWriteConfig writeConfig) throws IOException {
+    private void initImages(HoodieCDCFileSplit fileSplit) throws IOException {
       if (fileSplit.getBeforeFileSlice().isPresent() && !fileSplit.getBeforeFileSlice().get().isEmpty()) {
         this.beforeImages = imageManager.getOrLoadImages(
             maxCompactionMemoryInBytes, fileSplit.getBeforeFileSlice().get());
       } else {
-        this.beforeImages = FormatUtils.spillableMap(writeConfig, maxCompactionMemoryInBytes, getClass().getSimpleName());
+        this.beforeImages = Collections.emptyMap();
       }
     }
 
@@ -348,7 +350,6 @@ public RowData next() {
     @Override
     public void close() {
       logRecordIterator.close();
-      imageManager.close();
     }
 
     @SuppressWarnings("unchecked")
@@ -662,7 +663,12 @@ public BeforeImageIterator(
       this.maxCompactionMemoryInBytes = maxCompactionMemoryInBytes;
       this.projection = RowDataProjection.instance(requiredRowType, requiredPositions);
       this.imageManager = imageManager;
-      initImages(fileSplit);
+      try {
+        initImages(fileSplit);
+      } catch (IOException | RuntimeException | Error e) {
+        closeSuppressing(this, e);
+        throw e;
+      }
     }
 
     protected void initImages(HoodieCDCFileSplit fileSplit) throws IOException {