JVM：如何管理JNI创建的堆外内存

Question

我正在围绕Torch库构建一个Scala包装器。 我正在使用Swig构建胶层。 它允许我在堆外创建张量，我只能通过显式调用库的静态方法来释放它。 但是，我想以命令的方式使用张量，而不必担心释放内存，就像Java中的任何普通对象一样。

我能想到这样做的唯一方法是以下列方式（错误地）使用JVM的垃圾收集器：

“内存管理器”跟踪消耗的堆外内存量，当达到阈值时，它会调用System.gc() 。

object MemoryManager {    
  val Threshold: Long = 2L * 1024L * 1024L * 1024L // 2 GB
  val FloatSize = 4
  private val hiMemMark = new AtomicLong(0)

  def dec(size: Long): Long = hiMemMark.addAndGet(-size * FloatSize)
  def inc(size: Long): Long = hiMemMark.addAndGet(size * FloatSize)

  def memCheck(size: Long): Unit = {
    val level = inc(size)
    if (level > Threshold) {
      System.gc()
    }
  }
}

张量器本身包含在一个类中，使用finalize方法释放堆外内存，如下所示：

class Tensor private (val payload: SWIGTYPE_p_THFloatTensor) {
  def numel: Int = TH.THFloatTensor_numel(payload)

  override def finalize(): Unit = {
    val memSize = MemoryManager.dec(numel)
    TH.THFloatTensor_free(payload)
  }    
}

张量创建由工厂方法完成，该方法通知内存管理器。 例如，要创建零的Tensor：

object Tensor {
  def zeros(shape: List[Int]): Tensor = {
      MemoryManager.memCheck(shape.product)
      val storage = ... // boilerplate
      val t = TH.THFloatTensor_new
      TH.THFloatTensor_zeros(t, storage)
      new Tensor(t)
  }
}

我意识到这是一种天真的方法，但我可以逃脱这个吗？ 它似乎工作正常，也是在并行运行时（它会产生大量对System.gc()的多余调用，但除此之外什么都没有）或者你能想到一个更好的解决方案吗？

谢谢。

Answer 1

有一个更确定的选项 - 明确管理的内存区域

所以，大致如果我们有这样一个类：

class Region private () {
  private val registered = ArrayBuffer.empty[() => Unit]
  def register(finalizer: () => Unit): Unit = registered += finalizer
  def releaseAll(): Unit = {
    registered.foreach(f => f()) // todo - will leak if f() throws
  }
}

我们可以有一个实现所谓的“贷款模式”的方法，它给我们一个新的区域，然后处理释放

object Region {
  def run[A](f: Region => A): A = {
    val r = new Region
    try f(r) finally r.releaseAll()
  }
}

然后需要手动释放的东西可以被描述为采用隐式Region ：

class Leakable(i: Int)(implicit r: Region) {
  // Class body is constructor body, so you can register finalizers
  r.register(() => println(s"Deallocated foo $i"))

  def foo() = println(s"Foo: $i")
}

您可以以相当无样板的方式使用它：

Region.run { implicit r =>
  val a = new Leakable(1)
  val b = new Leakable(2)
  b.foo()
  a.foo()
}

此代码生成以下输出：

Foo: 2
Foo: 1
Deallocated foo 1
Deallocated foo 2

这种方法有一点限制（如果你尝试将Leakable分配给在run传递的闭包之外的变量，它的范围将不会被提升），但是即使禁用对System.gc调用，它也会更快并保证可以工作。