Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix JNI leak of a cudf::column_view native class. #10171

Merged
merged 1 commit into from
Feb 1, 2022
Merged

Fix JNI leak of a cudf::column_view native class. #10171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 24 additions & 16 deletions java/src/main/java/ai/rapids/cudf/ColumnVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@ public final class ColumnVector extends ColumnView {
NativeDepsLoader.loadNativeDeps();
}

private final OffHeapState offHeap;
private Optional<Long> nullCount = Optional.empty();
private int refCount;

Expand All @@ -56,14 +55,23 @@ public final class ColumnVector extends ColumnView {
* owned by this instance.
*/
public ColumnVector(long nativePointer) {
super(getColumnViewFromColumn(nativePointer));
super(new OffHeapState(nativePointer));
assert nativePointer != 0;
offHeap = new OffHeapState(nativePointer);
MemoryCleaner.register(this, offHeap);
this.refCount = 0;
incRefCountInternal(true);
}

private static OffHeapState makeOffHeap(DType type, long rows, Optional<Long> nullCount,
DeviceMemoryBuffer dataBuffer, DeviceMemoryBuffer validityBuffer,
DeviceMemoryBuffer offsetBuffer) {
long viewHandle = initViewHandle(
type, (int)rows, nullCount.orElse(UNKNOWN_NULL_COUNT).intValue(),
dataBuffer, validityBuffer, offsetBuffer, null);
return new OffHeapState(type, (int) rows, dataBuffer, validityBuffer,
offsetBuffer, null, viewHandle);
}

/**
* Create a new column vector based off of data already on the device.
* @param type the type of the vector
Expand All @@ -81,24 +89,29 @@ public ColumnVector(long nativePointer) {
public ColumnVector(DType type, long rows, Optional<Long> nullCount,
DeviceMemoryBuffer dataBuffer, DeviceMemoryBuffer validityBuffer,
DeviceMemoryBuffer offsetBuffer) {
super(ColumnVector.initViewHandle(
type, (int)rows, nullCount.orElse(UNKNOWN_NULL_COUNT).intValue(),
dataBuffer, validityBuffer, offsetBuffer, null));
super(makeOffHeap(type, rows, nullCount, dataBuffer, validityBuffer, offsetBuffer));
assert !type.equals(DType.LIST) : "This constructor should not be used for list type";
if (!type.equals(DType.STRING)) {
assert offsetBuffer == null : "offsets are only supported for STRING";
}
assert (nullCount.isPresent() && nullCount.get() <= Integer.MAX_VALUE)
|| !nullCount.isPresent();
offHeap = new OffHeapState(type, (int) rows, dataBuffer, validityBuffer,
offsetBuffer, null, viewHandle);
MemoryCleaner.register(this, offHeap);
this.nullCount = nullCount;

this.refCount = 0;
incRefCountInternal(true);
}

private static OffHeapState makeOffHeap(DType type, long rows, Optional<Long> nullCount,
DeviceMemoryBuffer dataBuffer, DeviceMemoryBuffer validityBuffer,
DeviceMemoryBuffer offsetBuffer, List<DeviceMemoryBuffer> toClose, long[] childHandles) {
long viewHandle = initViewHandle(type, (int)rows, nullCount.orElse(UNKNOWN_NULL_COUNT).intValue(),
dataBuffer, validityBuffer,
offsetBuffer, childHandles);
return new OffHeapState(type, (int) rows, dataBuffer, validityBuffer, offsetBuffer,
toClose, viewHandle);
}

/**
* Create a new column vector based off of data already on the device with child columns.
* @param type the type of the vector, typically a nested type
Expand All @@ -118,16 +131,12 @@ public ColumnVector(DType type, long rows, Optional<Long> nullCount,
public ColumnVector(DType type, long rows, Optional<Long> nullCount,
DeviceMemoryBuffer dataBuffer, DeviceMemoryBuffer validityBuffer,
DeviceMemoryBuffer offsetBuffer, List<DeviceMemoryBuffer> toClose, long[] childHandles) {
super(initViewHandle(type, (int)rows, nullCount.orElse(UNKNOWN_NULL_COUNT).intValue(),
dataBuffer, validityBuffer,
offsetBuffer, childHandles));
super(makeOffHeap(type, rows, nullCount, dataBuffer, validityBuffer, offsetBuffer, toClose, childHandles));
if (!type.equals(DType.STRING) && !type.equals(DType.LIST)) {
assert offsetBuffer == null : "offsets are only supported for STRING, LISTS";
}
assert (nullCount.isPresent() && nullCount.get() <= Integer.MAX_VALUE)
|| !nullCount.isPresent();
offHeap = new OffHeapState(type, (int) rows, dataBuffer, validityBuffer, offsetBuffer,
toClose, viewHandle);
MemoryCleaner.register(this, offHeap);

this.refCount = 0;
Expand All @@ -143,8 +152,7 @@ public ColumnVector(DType type, long rows, Optional<Long> nullCount,
* @param contiguousBuffer the buffer that this is based off of.
*/
private ColumnVector(long viewAddress, DeviceMemoryBuffer contiguousBuffer) {
super(viewAddress);
offHeap = new OffHeapState(viewAddress, contiguousBuffer);
super(new OffHeapState(viewAddress, contiguousBuffer));
MemoryCleaner.register(this, offHeap);
// TODO we may want to ask for the null count anyways...
this.nullCount = Optional.empty();
Expand Down
22 changes: 21 additions & 1 deletion java/src/main/java/ai/rapids/cudf/ColumnView.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ public class ColumnView implements AutoCloseable, BinaryOperable {
protected final DType type;
protected final long rows;
protected final long nullCount;
protected final ColumnVector.OffHeapState offHeap;

/**
* Constructs a Column View given a native view address
Expand All @@ -50,6 +51,22 @@ public class ColumnView implements AutoCloseable, BinaryOperable {
this.type = DType.fromNative(ColumnView.getNativeTypeId(viewHandle), ColumnView.getNativeTypeScale(viewHandle));
this.rows = ColumnView.getNativeRowCount(viewHandle);
this.nullCount = ColumnView.getNativeNullCount(viewHandle);
this.offHeap = null;
}


/**
* Intended to be called from ColumnVector when it is being constructed. Because state creates a
* cudf::column_view instance and will close it in all cases, we don't want to have to double
* close it.
* @param state the state this view is based off of.
*/
protected ColumnView(ColumnVector.OffHeapState state) {
offHeap = state;
viewHandle = state.getViewHandle();
type = DType.fromNative(ColumnView.getNativeTypeId(viewHandle), ColumnView.getNativeTypeScale(viewHandle));
rows = ColumnView.getNativeRowCount(viewHandle);
nullCount = ColumnView.getNativeNullCount(viewHandle);
}

/**
Expand Down Expand Up @@ -265,7 +282,10 @@ public long getDeviceMemorySize() {

@Override
public void close() {
ColumnView.deleteColumnView(viewHandle);
// close the view handle so long as offHeap is not going to do it for us.
if (offHeap == null) {
ColumnView.deleteColumnView(viewHandle);
}
viewHandle = 0;
}

Expand Down