NativeResource Management in Scala

scala-package/core/pom.xml

@@ -123,5 +123,12 @@
       <artifactId>commons-io</artifactId>
       <version>2.1</version>
     </dependency>
+    <!-- https://mvnrepository.com/artifact/org.mockito/mockito-all -->
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-all</artifactId>
+      <version>1.10.19</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 </project>

scala-package/core/src/main/scala/org/apache/mxnet/Executor.scala

@@ -45,7 +45,7 @@ object Executor {
  * @see Symbol.bind : to create executor
  */
 class Executor private[mxnet](private[mxnet] val handle: ExecutorHandle,
-                              private[mxnet] val symbol: Symbol) extends WarnIfNotDisposed {
+                              private[mxnet] val symbol: Symbol) extends NativeResource {
   private[mxnet] var argArrays: Array[NDArray] = null
   private[mxnet] var gradArrays: Array[NDArray] = null
   private[mxnet] var auxArrays: Array[NDArray] = null
@@ -59,16 +59,11 @@ class Executor private[mxnet](private[mxnet] val handle: ExecutorHandle,
   private[mxnet] var _group2ctx: Map[String, Context] = null
   private val logger: Logger = LoggerFactory.getLogger(classOf[Executor])
 
-  private var disposed = false
-  protected def isDisposed = disposed
-
-  def dispose(): Unit = {
-    if (!disposed) {
-      outputs.foreach(_.dispose())
-      _LIB.mxExecutorFree(handle)
-      disposed = true
-    }
-  }
+  override def nativeAddress: CPtrAddress = handle
+  override def nativeDeAllocator: (CPtrAddress => Int) = _LIB.mxExecutorFree
+  // cannot determine the off-heap size of this object
+  override val bytesAllocated: Long = 0
+  override val ref: NativeResourceRef = super.register()
 
   /**
    * Return a new executor with the same symbol and shared memory,
@@ -305,4 +300,5 @@ class Executor private[mxnet](private[mxnet] val handle: ExecutorHandle,
     checkCall(_LIB.mxExecutorPrint(handle, str))
     str.value
   }
+
 }

scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala

@@ -562,16 +562,20 @@ object NDArray extends NDArrayBase {
  */
 class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle,
                              val writable: Boolean = true,
-                             addToCollector: Boolean = true) extends WarnIfNotDisposed {
+                             addToCollector: Boolean = true) extends NativeResource {
   if (addToCollector) {
     NDArrayCollector.collect(this)
   }
 
+  override def nativeAddress: CPtrAddress = handle
+  override def nativeDeAllocator: (CPtrAddress => Int) = _LIB.mxNDArrayFree
+  override val bytesAllocated: Long = DType.numOfBytes(this.dtype) * this.shape.product
+
+  override val ref: NativeResourceRef = super.register()
+
   // record arrays who construct this array instance
   // we use weak reference to prevent gc blocking
   private[mxnet] val dependencies = mutable.HashMap.empty[Long, WeakReference[NDArray]]
-  @volatile private var disposed = false
-  def isDisposed: Boolean = disposed
 
   def serialize(): Array[Byte] = {
     val buf = ArrayBuffer.empty[Byte]
@@ -584,11 +588,10 @@ class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle,
    * The NDArrays it depends on will NOT be disposed. <br />
    * The object shall never be used after it is disposed.
    */
-  def dispose(): Unit = {
+  override def dispose(): Unit = {
     if (!disposed) {
-      _LIB.mxNDArrayFree(handle)
+      super.dispose()
       dependencies.clear()
-      disposed = true
     }
   }
 
@@ -1034,6 +1037,7 @@ class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle,
     // TODO: naive implementation
     shape.hashCode + toArray.hashCode
   }
+
 }
 
 private[mxnet] object NDArrayConversions {

scala-package/core/src/main/scala/org/apache/mxnet/NativeResource.scala

@@ -0,0 +1,186 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.mxnet
+
+import org.apache.mxnet.Base.CPtrAddress
+import java.lang.ref.{PhantomReference, ReferenceQueue, WeakReference}
+import java.util.concurrent._
+
+import org.apache.mxnet.Base.checkCall
+import java.util.concurrent.atomic.AtomicLong
+
+
+/**
+  * NativeResource trait is used to manage MXNet Objects
+  * such as NDArray, Symbol, Executor, etc.,
+  * The MXNet Object calls NativeResource.register
+  * and assign the returned NativeResourceRef to PhantomReference
+  * NativeResource also implements AutoCloseable so MXNetObjects
+  * can be used like Resources in try-with-resources paradigm
+  */
+private[mxnet] trait NativeResource
+  extends AutoCloseable with WarnIfNotDisposed {
+
+  /**
+    * native Address associated with this object
+    */
+  def nativeAddress: CPtrAddress
+
+  /**
+    * Function Pointer to the NativeDeAllocator of nativeAddress
+    */
+  def nativeDeAllocator: (CPtrAddress => Int)
+
+  /** Call NativeResource.register to get the reference
+    */
+  val ref: NativeResourceRef
+
+  /**
+    * Off-Heap Bytes Allocated for this object
+    */
+  // intentionally making it a val, so it gets evaluated when defined
+  val bytesAllocated: Long
+
+  private[mxnet] var scope: ResourceScope = null
+
+  @volatile var disposed = false
+
+  override def isDisposed: Boolean = disposed || isDeAllocated
+
+  /**
+    * Register this object for PhantomReference tracking and in
+    * ResourceScope if used inside ResourceScope.
+    * @return NativeResourceRef that tracks reachability of this object
+    *         using PhantomReference
+    */
+  def register(): NativeResourceRef = {
+    scope = ResourceScope.getCurrentScope()
+    if (scope != null) scope.add(this)
+
+    NativeResource.totalBytesAllocated.getAndAdd(bytesAllocated)
+    // register with PhantomRef tracking to release incase the objects go
+    // out of reference within scope but are held for long time
+    NativeResourceRef.register(this, nativeDeAllocator)
+ }
+
+  // Removes this object from PhantomRef tracking and from ResourceScope
+  private def deRegister(removeFromScope: Boolean): Unit = {
+    NativeResourceRef.deRegister(ref)
+    if (scope != null && removeFromScope) scope.remove(this)
+  }
+
+  // Implements [[@link AutoCloseable.close]]
+  override def close(): Unit = {
+    dispose()
+  }
+
+  // Implements [[@link WarnIfNotDisposed.dispose]]
+  def dispose(): Unit = dispose(true)
+
+  private[mxnet] def dispose(removeFromScope: Boolean = true): Unit = {
+    if (!disposed) {
+      checkCall(nativeDeAllocator(this.nativeAddress))
+      deRegister(removeFromScope)
+      NativeResource.totalBytesAllocated.getAndAdd(-1*bytesAllocated)
+      disposed = true
+    }
+  }
+
+  /*
+  this is used by the WarnIfNotDisposed finalizer,
+  the object could be disposed by the GC without the need for explicit disposal
+  but the finalizer might not have run, then the WarnIfNotDisposed throws a warning
+   */
+  private[mxnet] def isDeAllocated(): Boolean = NativeResourceRef.isDeAllocated(ref)
+
+}
+
+private[mxnet] object NativeResource {
+  var totalBytesAllocated : AtomicLong = new AtomicLong(0)
+}
+
+// Do not make [[NativeResource.resource]] a member of the class,
+// this will hold reference and GC will not clear the object.
+private[mxnet] class NativeResourceRef(resource: NativeResource,
+                                       val resourceDeAllocator: CPtrAddress => Int)
+        extends PhantomReference[NativeResource](resource, NativeResourceRef.refQ) {}
+
+private[mxnet] object NativeResourceRef {
+
+  private[mxnet] val refQ: ReferenceQueue[NativeResource]
+                = new ReferenceQueue[NativeResource]
+
+  private[mxnet] val refMap = new ConcurrentHashMap[NativeResourceRef, CPtrAddress]()
+
+  private[mxnet] val cleaner = new ResourceCleanupThread()
+
+  cleaner.start()
+
+  def register(resource: NativeResource, nativeDeAllocator: (CPtrAddress => Int)):
+  NativeResourceRef = {
+    val ref = new NativeResourceRef(resource, nativeDeAllocator)
+    refMap.put(ref, resource.nativeAddress)
+    ref
+  }
+
+  def deRegister(ref: NativeResourceRef): Unit = {
+    if (refMap.containsKey(ref)) {
+      refMap.remove(ref)
+    }
+  }
+
+  /**
+    * This method will check if the cleaner ran and deAllocated the object
+    * As a part of GC, when the object is unreachable GC inserts a phantomRef
+    * to the ReferenceQueue which the cleaner thread will deallocate, however
+    * the finalizer runs much later depending on the GC.
+    * @param resource resource to verify if it has been deAllocated
+    * @return true if already deAllocated
+    */
+  def isDeAllocated(ref: NativeResourceRef): Boolean = {
+    !refMap.containsKey(ref)
+  }
+
+  def cleanup: Unit = {
+    // remove is a blocking call
+    val ref: NativeResourceRef = refQ.remove().asInstanceOf[NativeResourceRef]
+    // phantomRef will be removed from the map when NativeResource.close is called.
+    val resource = refMap.get(ref)
+    if (resource != 0L)  { // since CPtrAddress is Scala a Long, it cannot be null
+      ref.resourceDeAllocator(resource)
+      refMap.remove(ref)
+    }
+  }
+
+  protected class ResourceCleanupThread extends Thread {
+    setPriority(Thread.MAX_PRIORITY)
+    setName("NativeResourceDeAllocatorThread")
+    setDaemon(true)
+
+    override def run(): Unit = {
+      while (true) {
+        try {
+          NativeResourceRef.cleanup
+        }
+        catch {
+          case _: InterruptedException => Thread.currentThread().interrupt()
+        }
+      }
+    }
+  }
+}