java - 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)
  }    
}

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

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 的调用，速度也会更快并保证正常工作。