From 27acfa45630953a17c2221420a1c45597afdf5a5 Mon Sep 17 00:00:00 2001
From: Jason Lowe <jlowe@nvidia.com>
Date: Mon, 10 Aug 2020 20:55:49 -0500
Subject: [PATCH 1/2] Fix reference count on Java DeviceMemoryBuffer after
 contiguousSplit

---
 .../main/java/ai/rapids/cudf/ColumnVector.java  |  2 ++
 .../java/ai/rapids/cudf/DeviceMemoryBuffer.java |  6 ++----
 .../main/java/ai/rapids/cudf/MemoryBuffer.java  | 17 +++++++++++++++--
 .../java/ai/rapids/cudf/ColumnVectorTest.java   |  9 +++++++++
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/java/src/main/java/ai/rapids/cudf/ColumnVector.java b/java/src/main/java/ai/rapids/cudf/ColumnVector.java
index c49169bff79..f486fa2d115 100644
--- a/java/src/main/java/ai/rapids/cudf/ColumnVector.java
+++ b/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() {
diff --git a/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java
index a2e95d26fa7..e3184e311f4 100644
--- a/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java
+++ b/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);
   }
 }
diff --git a/java/src/main/java/ai/rapids/cudf/MemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/MemoryBuffer.java
index 45b352d3604..60a69881491 100644
--- a/java/src/main/java/ai/rapids/cudf/MemoryBuffer.java
+++ b/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;
+  }
 }
diff --git a/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java b/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
index bca8e0baf93..4fc96f8b065 100644
--- a/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
+++ b/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());
+    }
+  }
 }

From 41ce1e904884a121010b810c6a52f690c1898b3d Mon Sep 17 00:00:00 2001
From: Jason Lowe <jlowe@nvidia.com>
Date: Mon, 10 Aug 2020 21:09:50 -0500
Subject: [PATCH 2/2] changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b78a9c74d11..9c960f52bd5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -275,6 +275,7 @@
 - PR #5872 libcudf_kafka r_path is causing docker build failures on centos7
 - PR #5869 Fix bug in parquet writer in writing string column with offset
 - 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)