Java cleaner synchronization (rapidsai#7474)

Add synchronization in `cleanImpl` and `close` in various places where race conditions could exist, and also within the `MemoryCleaner` to address some concurrent modification issues we've seen in tests while shutting down (i.e. invoking the cleaner) (i.e. NVIDIA/spark-rapids#1797) Authors: - Alessandro Bellina (@abellina) Approvers: - Robert (Bobby) Evans (@revans2) - Jason Lowe (@jlowe) URL: rapidsai#7474
hyperbolic2346 · Mar 23, 2021 · cc73af6 · cc73af6
1 parent 76894e1
commit cc73af6
Show file tree

Hide file tree

Showing 11 changed files with 26 additions and 20 deletions.
diff --git a/java/src/main/java/ai/rapids/cudf/ColumnVector.java b/java/src/main/java/ai/rapids/cudf/ColumnVector.java
@@ -858,7 +858,7 @@ public String toString() {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long address = 0;
 

diff --git a/java/src/main/java/ai/rapids/cudf/Cuda.java b/java/src/main/java/ai/rapids/cudf/Cuda.java
@@ -41,7 +41,7 @@ private static class StreamCleaner extends MemoryCleaner.Cleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long origAddress = stream;
       if (stream != CUDA_STREAM_DEFAULT &&
@@ -114,7 +114,7 @@ public String toString() {
     }
 
     @Override
-    public void close() {
+    public synchronized void close() {
       if (cleaner != null) {
         cleaner.delRef();
       }
@@ -139,7 +139,7 @@ private static class EventCleaner extends MemoryCleaner.Cleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long origAddress = event;
       if (event != 0) {
@@ -233,7 +233,7 @@ public String toString() {
     }
 
     @Override
-    public void close() {
+    public synchronized void close() {
       cleaner.delRef();
       if (closed) {
         cleaner.logRefCountDebug("double free " + this);

diff --git a/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java
@@ -41,7 +41,7 @@ private static final class DeviceBufferCleaner extends MemoryBufferCleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long origAddress = address;
       if (address != 0) {
@@ -79,7 +79,7 @@ private static final class RmmDeviceBufferCleaner extends MemoryBufferCleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       if (rmmBufferAddress != 0) {
         Rmm.freeDeviceBuffer(rmmBufferAddress);

diff --git a/java/src/main/java/ai/rapids/cudf/HostColumnVectorCore.java b/java/src/main/java/ai/rapids/cudf/HostColumnVectorCore.java
@@ -514,7 +514,7 @@ public long getHostMemorySize() {
    * Close method for the column
    */
   @Override
-  public void close() {
+  public synchronized void close() {
     for (HostColumnVectorCore child : children) {
       if (child != null) {
         child.close();
@@ -553,7 +553,7 @@ protected static final class OffHeapState extends MemoryCleaner.Cleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       if (data != null || valid != null || offsets != null) {
         try {

diff --git a/java/src/main/java/ai/rapids/cudf/HostMemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/HostMemoryBuffer.java
@@ -63,7 +63,7 @@ private static final class HostBufferCleaner extends MemoryBufferCleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long origAddress = address;
       if (address != 0) {

diff --git a/java/src/main/java/ai/rapids/cudf/MemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/MemoryBuffer.java
@@ -46,7 +46,7 @@ private static final class SlicedBufferCleaner extends MemoryBufferCleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       if (parent != null) {
         if (logErrorIfNotClean) {
           log.error("A SLICED BUFFER WAS LEAKED(ID: " + id + " parent: " + parent + ")");

diff --git a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java b/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java
@@ -79,19 +79,25 @@ public Cleaner() {
 
     public final void addRef() {
       if (REF_COUNT_DEBUG && refCountDebug != null) {
-        refCountDebug.add(new MemoryCleaner.RefCountDebugItem("INC"));
+        synchronized(this)  {
+          refCountDebug.add(new MemoryCleaner.RefCountDebugItem("INC"));
+        }
       }
     }
 
     public final void delRef() {
       if (REF_COUNT_DEBUG && refCountDebug != null) {
-        refCountDebug.add(new MemoryCleaner.RefCountDebugItem("DEC"));
+        synchronized(this) {
+          refCountDebug.add(new MemoryCleaner.RefCountDebugItem("DEC"));
+        }
       }
     }
 
     public final void logRefCountDebug(String message) {
       if (REF_COUNT_DEBUG && refCountDebug != null) {
-        log.error("{} (ID: {}): {}", message, id, MemoryCleaner.stringJoin("\n", refCountDebug));
+        synchronized(this) {
+          log.error("{} (ID: {}): {}", message, id, MemoryCleaner.stringJoin("\n", refCountDebug));
+        }
       }
     }
 

diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java
@@ -108,7 +108,7 @@ private static final class PinnedHostBufferCleaner extends MemoryBuffer.MemoryBu
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long origAddress = 0;
       if (section != null) {

diff --git a/java/src/main/java/ai/rapids/cudf/Scalar.java b/java/src/main/java/ai/rapids/cudf/Scalar.java
@@ -675,7 +675,7 @@ private static class OffHeapState extends MemoryCleaner.Cleaner {
     }
 
     @Override
-    protected boolean cleanImpl(boolean logErrorIfNotClean) {
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
       if (scalarHandle != 0) {
         if (logErrorIfNotClean) {
           LOG.error("A SCALAR WAS LEAKED(ID: " + id + " " + Long.toHexString(scalarHandle) + ")");

diff --git a/java/src/main/java/ai/rapids/cudf/nvcomp/BatchedLZ4Decompressor.java b/java/src/main/java/ai/rapids/cudf/nvcomp/BatchedLZ4Decompressor.java
@@ -175,7 +175,7 @@ public boolean isLZ4Metadata() {
     }
 
     @Override
-    public void close() {
+    public synchronized void close() {
       if (!closed) {
         cleaner.delRef();
         cleaner.clean(false);
@@ -200,7 +200,7 @@ private static class BatchedMetadataCleaner extends MemoryCleaner.Cleaner {
       }
 
       @Override
-      protected boolean cleanImpl(boolean logErrorIfNotClean) {
+      protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
         boolean neededCleanup = false;
         long address = metadata;
         if (metadata != 0) {

diff --git a/java/src/main/java/ai/rapids/cudf/nvcomp/Decompressor.java b/java/src/main/java/ai/rapids/cudf/nvcomp/Decompressor.java
@@ -110,7 +110,7 @@ public boolean isLZ4Metadata() {
     }
 
     @Override
-    public void close() {
+    public synchronized void close() {
       if (!closed) {
         cleaner.delRef();
         cleaner.clean(false);
@@ -135,7 +135,7 @@ private static class MetadataCleaner extends MemoryCleaner.Cleaner {
       }
 
       @Override
-      protected boolean cleanImpl(boolean logErrorIfNotClean) {
+      protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
         boolean neededCleanup = false;
         long address = metadata;
         if (metadata != 0) {