Fixed a memory leak caused by Exception thrown while constructing a C…

…olumnView (#13597)

Close unclosed buffers in case of any Exception that is thrown in the ColumnView constructor. The catch all Exception clause then caused a double-free corruption which was also fixed

Authors:
  - Raza Jafri (https://github.com/razajafri)

Approvers:
  - Jason Lowe (https://github.com/jlowe)

URL: #13597

java/src/main/java/ai/rapids/cudf/ColumnVector.java

@@ -1728,4 +1728,39 @@ public static ColumnVector empty(HostColumnVector.DataType colType) {
       throw new IllegalArgumentException("Unsupported data type: " + colType);
     }
   }
+
+  static ColumnVector[] getColumnVectorsFromPointers(long[] nativeHandles) {
+    ColumnVector[] columns = new ColumnVector[nativeHandles.length];
+    try {
+      for (int i = 0; i < nativeHandles.length; i++) {
+        long nativeHandle = nativeHandles[i];
+        // setting address to zero, so we don't clean it in case of an exception as it
+        // will be cleaned up by the constructor
+        nativeHandles[i] = 0;
+        columns[i] = new ColumnVector(nativeHandle);
+      }
+      return columns;
+    } catch (Throwable t) {
+      for (ColumnVector columnVector : columns) {
+        if (columnVector != null) {
+          try {
+            columnVector.close();
+          } catch (Throwable s) {
+            t.addSuppressed(s);
+          }
+        }
+      }
+      for (long nativeHandle : nativeHandles) {
+        if (nativeHandle != 0) {
+          try {
+            deleteCudfColumn(nativeHandle);
+          } catch (Throwable s) {
+            t.addSuppressed(s);
+          }
+        }
+      }
+
+      throw t;
+    }
+  }
 }

java/src/main/java/ai/rapids/cudf/ColumnView.java

@@ -56,12 +56,12 @@ public class ColumnView implements AutoCloseable, BinaryOperable {
     this.offHeap = null;
     try {
       AssertEmptyNulls.assertNullsAreEmpty(this);
-    } catch (AssertionError ae) {
+    } catch (Throwable t) {
       // offHeap state is null, so there is nothing to clean in offHeap
       // delete ColumnView to avoid memory leak
       deleteColumnView(viewHandle);
       viewHandle = 0;
-      throw ae;
+      throw t;
     }
   }
 
@@ -81,11 +81,11 @@ protected ColumnView(ColumnVector.OffHeapState state) {
     nullCount = ColumnView.getNativeNullCount(viewHandle);
     try {
       AssertEmptyNulls.assertNullsAreEmpty(this);
-    } catch (AssertionError ae) {
+    } catch (Throwable t) {
       // cleanup offHeap
       offHeap.clean(false);
       viewHandle = 0;
-      throw ae;
+      throw t;
     }
   }
 
@@ -673,8 +673,13 @@ public final ColumnVector[] slice(int... indices) {
         nativeHandles[i] = 0;
       }
     } catch (Throwable t) {
-      cleanupColumnViews(nativeHandles, columnVectors);
-      throw t;
+      try {
+        cleanupColumnViews(nativeHandles, columnVectors, t);
+      } catch (Throwable s) {
+        t.addSuppressed(s);
+      } finally {
+        throw t;
+      }
     }
     return columnVectors;
   }
@@ -818,21 +823,34 @@ public ColumnView[] splitAsViews(int... indices) {
         nativeHandles[i] = 0;
       }
     } catch (Throwable t) {
-      cleanupColumnViews(nativeHandles, columnViews);
-      throw t;
+      try {
+        cleanupColumnViews(nativeHandles, columnViews, t);
+      } catch (Throwable s) {
+        t.addSuppressed(s);
+      } finally {
+        throw t;
+      }
     }
     return columnViews;
   }
 
-  static void cleanupColumnViews(long[] nativeHandles, ColumnView[] columnViews) {
-    for (ColumnView columnView: columnViews) {
+  static void cleanupColumnViews(long[] nativeHandles, ColumnView[] columnViews, Throwable throwable) {
+    for (ColumnView columnView : columnViews) {
       if (columnView != null) {
-        columnView.close();
+        try {
+          columnView.close();
+        } catch (Throwable s) {
+          throwable.addSuppressed(s);
+        }
       }
     }
-    for (long nativeHandle: nativeHandles) {
+    for (long nativeHandle : nativeHandles) {
       if (nativeHandle != 0) {
-        deleteColumnView(nativeHandle);
+        try {
+          deleteColumnView(nativeHandle);
+        } catch (Throwable s) {
+          throwable.addSuppressed(s);
+        }
       }
     }
   }
@@ -5168,4 +5186,26 @@ public boolean hasNonEmptyNulls() {
   public ColumnVector purgeNonEmptyNulls() {
     return new ColumnVector(purgeNonEmptyNulls(viewHandle));
   }
+
+  static ColumnView[] getColumnViewsFromPointers(long[] nativeHandles) {
+    ColumnView[] columns = new ColumnView[nativeHandles.length];
+    try {
+      for (int i = 0; i < nativeHandles.length; i++) {
+        long nativeHandle = nativeHandles[i];
+        // setting address to zero, so we don't clean it in case of an exception as it
+        // will be cleaned up by the constructor
+        nativeHandles[i] = 0;
+        columns[i] = new ColumnView(nativeHandle);
+      }
+      return columns;
+    } catch (Throwable t) {
+      try {
+        cleanupColumnViews(nativeHandles, columns, t);
+      } catch (Throwable s) {
+        t.addSuppressed(s);
+      } finally {
+        throw t;
+      }
+    }
+  }
 }

java/src/main/java/ai/rapids/cudf/Scalar.java

@@ -677,26 +677,7 @@ public ColumnView[] getChildrenFromStructScalar() {
     assert DType.STRUCT.equals(type) : "Cannot get table for the vector of type " + type;
 
     long[] childHandles = getChildrenFromStructScalar(getScalarHandle());
-    ColumnView[] children = new ColumnView[childHandles.length];
-    try {
-      for (int i = 0; i < children.length; i++) {
-        children[i] = new ColumnView(childHandles[i]);
-      }
-    } catch (Exception ex) {
-      // close all created ColumnViews if exception thrown
-      for (ColumnView child : children) {
-        // We closed all created ColumnViews when we hit null. Therefore we exit the loop.
-        if (child == null) break;
-        // make sure the close process is exception-free
-        try {
-          child.close();
-        } catch (Exception suppressed) {
-          ex.addSuppressed(suppressed);
-        }
-      }
-      throw ex;
-    }
-    return children;
+    return ColumnView.getColumnViewsFromPointers(childHandles);
   }
 
   @Override

java/src/main/java/ai/rapids/cudf/Table.java

@@ -83,20 +83,22 @@ public Table(ColumnVector... columns) {
    */
   public Table(long[] cudfColumns) {
     assert cudfColumns != null && cudfColumns.length > 0 : "CudfColumns can't be null or empty";
-    this.columns = new ColumnVector[cudfColumns.length];
+    this.columns = ColumnVector.getColumnVectorsFromPointers(cudfColumns);
     try {
-      for (int i = 0; i < cudfColumns.length; i++) {
-        this.columns[i] = new ColumnVector(cudfColumns[i]);
-        cudfColumns[i] = 0;
-      }
       long[] views = new long[columns.length];
       for (int i = 0; i < columns.length; i++) {
         views[i] = columns[i].getNativeView();
       }
       nativeHandle = createCudfTableView(views);
       this.rows = columns[0].getRowCount();
     } catch (Throwable t) {
-      ColumnView.cleanupColumnViews(cudfColumns, this.columns);
+      for (ColumnVector column : columns) {
+        try {
+          column.close();
+        } catch (Throwable s) {
+          t.addSuppressed(s);
+        }
+      }
       throw t;
     }
   }
@@ -3462,17 +3464,7 @@ public static GatherMap mixedLeftAntiJoinGatherMap(Table leftKeys, Table rightKe
    */
   public ColumnVector[] convertToRows() {
     long[] ptrs = convertToRows(nativeHandle);
-    ColumnVector[] ret = new ColumnVector[ptrs.length];
-    try {
-      for (int i = 0; i < ptrs.length; i++) {
-        ret[i] = new ColumnVector(ptrs[i]);
-        ptrs[i] = 0;
-      }
-    } catch (Throwable t) {
-      ColumnView.cleanupColumnViews(ptrs, ret);
-      throw t;
-    }
-    return ret;
+    return ColumnVector.getColumnVectorsFromPointers(ptrs);
   }
 
   /**
@@ -3551,17 +3543,7 @@ public ColumnVector[] convertToRows() {
    */
   public ColumnVector[] convertToRowsFixedWidthOptimized() {
     long[] ptrs = convertToRowsFixedWidthOptimized(nativeHandle);
-    ColumnVector[] ret = new ColumnVector[ptrs.length];
-    try {
-      for (int i = 0; i < ptrs.length; i++) {
-        ret[i] = new ColumnVector(ptrs[i]);
-        ptrs[i] = 0;
-      }
-    } catch (Throwable t) {
-      ColumnView.cleanupColumnViews(ptrs, ret);
-      throw t;
-    }
-    return ret;
+    return ColumnVector.getColumnVectorsFromPointers(ptrs);
   }
 
   /**
@@ -3626,13 +3608,21 @@ public static Table fromPackedTable(ByteBuffer metadata, DeviceMemoryBuffer data
     Table result = null;
     try {
       for (int i = 0; i < columns.length; i++) {
-        columns[i] = ColumnVector.fromViewWithContiguousAllocation(columnViewAddresses[i], data);
+        long columnViewAddress = columnViewAddresses[i];
+        // setting address to zero, so we don't clean it in case of an exception as it
+        // will be cleaned up by the ColumnView  constructor
         columnViewAddresses[i] = 0;
+        columns[i] = ColumnVector.fromViewWithContiguousAllocation(columnViewAddress, data);
       }
       result = new Table(columns);
     } catch (Throwable t) {
-      ColumnView.cleanupColumnViews(columnViewAddresses, columns);
-      throw t;
+      try {
+        ColumnView.cleanupColumnViews(columnViewAddresses, columns, t);
+      } catch (Throwable s){
+        t.addSuppressed(s);
+      } finally {
+        throw t;
+      }
     }
 
     // close columns to leave the resulting table responsible for freeing underlying columns