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

Question

我正在构建一个Scala封装器，围绕Torch库。我正在使用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)
  }    
}

张量的创建是通过工厂方法完成的，该方法会通知内存管理器。例如，要创建一个全为零的张量：

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() 调用，但除此之外没有其他问题）。或者您能想到更好的解决方案吗？

谢谢。

- botkop

---

这可能很好用。我的做法类似于使用直接 ByteBuffer 来完成。 - Peter Lawrey

@PeterLawrey 这很有趣，但我找不到他们在JVM代码或文档中如何实现的。你能提供一个链接吗？谢谢。 - botkop

我建议阅读源代码。您可以在IDE中单击它或使用http://grepcode.com/file/repository.grepcode.com/java/root/jdk/openjdk/8u40-b25/java/nio/DirectByteBuffer.java?av=f。 - Peter Lawrey

如果你查看它的父级，你会发现http://grepcode.com/file/repository.grepcode.com/java/root/jdk/openjdk/8u40-b25/java/nio/Buffer.java#Buffer - Peter Lawrey

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分配给闭包外部的变量，则其作用域不会被提升），但速度更快，并且即使禁用对System.gc的调用，也可以保证正常工作。

- Oleg Pyzhcov

---

谢谢Oleg，我会让这个想法沉淀一段时间。 - botkop

我认为这是一个非常优雅的解决方案。谢谢。 - botkop