Merge pull request #5915 from jlowe/jni-contigsplit-refcount

Fix reference count on Java DeviceMemoryBuffer after contiguousSplit

CHANGELOG.md

@@ -276,6 +276,7 @@
 - PR #5869 Fix bug in parquet writer in writing string column with offset
 - PR #5914 Link CUDA against libcudf_kafka
 - PR #5895 Do not break kafka client consumption loop on local client timeout
+- PR #5915 Fix reference count on Java DeviceMemoryBuffer after contiguousSplit
 
 
 # cuDF 0.14.0 (03 Jun 2020)

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

@@ -3079,6 +3079,8 @@ public OffHeapState(long viewHandle, DeviceMemoryBuffer contiguousBuffer) {
       toClose.add(data);
       toClose.add(valid);
       toClose.add(offsets);
+      contiguousBuffer.incRefCount();
+      toClose.add(contiguousBuffer);
     }
 
     public long getViewHandle() {

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

@@ -136,8 +136,7 @@ public static DeviceMemoryBuffer allocate(long bytes) {
   @Override
   public synchronized final DeviceMemoryBuffer slice(long offset, long len) {
     addressOutOfBoundsCheck(address + offset, len, "slice");
-    refCount++;
-    cleaner.addRef();
+    incRefCount();
     return new DeviceMemoryBuffer(getAddress() + offset, len, this);
   }
 
@@ -151,8 +150,7 @@ synchronized final BaseDeviceMemoryBuffer sliceFrom(DeviceMemoryBufferView view)
       return null;
     }
     addressOutOfBoundsCheck(view.address, view.length, "sliceFrom");
-    refCount++;
-    cleaner.addRef();
+    incRefCount();
     return new DeviceMemoryBuffer(view.address, view.length, this);
   }
 }

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

@@ -95,8 +95,7 @@ protected MemoryBuffer(long address, long length, MemoryBufferCleaner cleaner) {
     this.cleaner = cleaner;
     if (cleaner != null) {
       this.id = cleaner.id;
-      refCount++;
-      cleaner.addRef();
+      incRefCount();
       MemoryCleaner.register(this, cleaner);
     } else {
       this.id = -1;
@@ -190,4 +189,18 @@ public String toString() {
         ", length=" + length +
         ", id=" + id + "}";
   }
+
+  /**
+   * Increment the reference count for this column.  You need to call close on this
+   * to decrement the reference count again.
+   */
+  public synchronized void incRefCount() {
+    refCount++;
+    cleaner.addRef();
+  }
+
+  // visible for testing
+  synchronized int getRefCount() {
+    return refCount;
+  }
 }

java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java

@@ -2829,4 +2829,13 @@ void testConcatListsOfLists() {
       assertColumnsAreEqual(expected, v);
     }
   }
+
+  @Test
+  void testContiguousSplitConstructor() {
+    try (Table tmp = new Table.TestBuilder().column(1, 2).column(3, 4).build();
+         ContiguousTable ct = tmp.contiguousSplit()[0]) {
+      // one reference for the device buffer itself, two more for the column using it
+      assertEquals(3, ct.getBuffer().getRefCount());
+    }
+  }
 }