Kafka Streams Scala自定义SuppressTransformer内存泄漏问题排查求助

嘿，我看了你自定义实现壁钟时间Suppress功能后遇到的内存不释放问题，咱们一步步拆解核心问题和修复方案：

1. 最致命的问题：keySet 导致的内存泄漏 #

你代码里用了全局的mutable.HashSet[String] keySet来追踪需要处理的key，但这里有两个大坑：

• 一个session key可能对应多个窗口，你删除其中一个窗口就把key从keySet里移除了，剩下的窗口永远不会被处理；
• 如果某个key的所有窗口都处理完了，但后续没有新数据进来，这个key会一直留在keySet里，永久占用内存。

修复思路： #

直接抛弃手动维护keySet的逻辑，每次调度时直接遍历状态存储里的所有窗口key，这样不会遗漏任何过期窗口，也不会残留无效key：

this.context.schedule(
  scheduleInterval,
  PunctuationType.WALL_CLOCK_TIME,
  _ => {
    val allWindowedKeys = store.all() // 获取所有session窗口的完整迭代器
    val now = Instant.now()
    val expiredThreshold = now.minus(windowDuration).toEpochMilli

    try {
      while (allWindowedKeys.hasNext) {
        val entry = allWindowedKeys.next()
        val windowedKey = entry.key
        val windowStartTime = windowedKey.window().start()

        // 判断窗口是否过期
        if (windowStartTime > 0 && windowStartTime < expiredThreshold) {
          // 转发聚合数据到下游
          context.forward(windowedKey.key(), entry.value, To.all().withTimestamp(now.toEpochMilli))
          // 删除状态存储中的对应窗口
          store.remove(windowedKey)
        }
      }
    } finally {
      allWindowedKeys.close() // 必须关闭迭代器，避免资源泄漏
    }

    // 统一提交状态变更，不要在循环里频繁commit
    context.commit()
  }
)

同时删掉transform方法里所有和keySet相关的代码——现在不需要手动追踪key了。

2. RocksDB 的内存回收机制坑 #

你用RocksDB作为状态存储，它的删除操作是标记删除，不会立即释放内存，要等后台压缩和垃圾回收触发。另外默认的内存缓存配置可能偏大，导致内存迟迟不下降。

优化建议： #

配置RocksDB时，明确限制block cache大小并开启自动压缩：

new RocksDbSessionBytesStoreSupplier(
  stateStoreName,
  stateStoreRetention.toMillis,
  256 * 1024 * 1024, // 限制block cache为256MB，根据你的机器调整
  true // 开启自动压缩
)

测试时可以多等一会儿，或者查看RocksDB状态目录（默认在/tmp/kafka-streams/<你的应用ID>/state）的文件大小，确认数据确实被清理了。

3. 迭代器与commit时机的问题 #

原来的代码里，你在循环内调用context.commit()会导致频繁的状态提交，既影响性能，还可能导致数据不一致。而且用store.fetch(key)遍历单个key的窗口，容易遗漏窗口或者引发迭代器异常。

修复要点： #

• 用store.all()遍历所有窗口，确保不遗漏任何过期数据；
• 把context.commit()移到遍历完成后的finally块外，统一提交状态变更。

4. Session Window Grace Period 的冲突 #

你代码里配置了SessionWindows with sessionWindowMinDuration grace sessionGracePeriodDuration，如果grace period设置得比你的自定义windowDuration大，Kafka Streams可能会强制保留状态直到grace period结束。建议把sessionGracePeriodDuration设置成小于等于你的自定义窗口时长，或者在逻辑里忽略这个限制（根据业务需求调整）。

完整修改后的 SuppressTransformer 代码 #

class SuppressTransformer[T](stateStoreName: String, windowDuration: Duration) extends Transformer[String, T, KeyValue[String, T]] {
  val scheduleInterval: Duration = Duration.ofSeconds(180)
  var context: ProcessorContext = _
  var store: SessionStore[String, Array[T]] = _

  override def init(context: ProcessorContext): Unit = {
    this.context = context;
    this.store = context.getStateStore(stateStoreName).asInstanceOf[SessionStore[String, Array[T]]]
    this.context.schedule(
      scheduleInterval,
      PunctuationType.WALL_CLOCK_TIME,
      _ => {
        val allWindowedKeys = store.all()
        val now = Instant.now()
        val expiredWindowThreshold = now.minus(windowDuration).toEpochMilli

        try {
          while (allWindowedKeys.hasNext) {
            val entry = allWindowedKeys.next()
            val windowedKey = entry.key
            val windowStartTime = windowedKey.window().start()

            if (windowStartTime > 0 && windowStartTime < expiredWindowThreshold) {
              context.forward(windowedKey.key(), entry.value, To.all().withTimestamp(now.toEpochMilli))
              store.remove(windowedKey)
            }
          }
        } finally {
          allWindowedKeys.close() // 确保迭代器关闭，避免资源泄漏
        }

        context.commit()
      }
    )
  }

  override def transform(key: String, value: T): KeyValue[String, T] = {
    null // 不需要处理，直接让上游聚合结果进入状态存储
  }

  override def close(): Unit = {}
}

测试验证小技巧 #

• 用JVM内存分析工具（比如VisualVM）查看堆内存对象，确认没有残留的keySet或未关闭的迭代器；
• 通过Kafka Streams的InteractiveQuery API查询状态存储，确认过期窗口确实被删除；
• 观察RocksDB状态目录的文件大小，确认数据被清理后文件体积下降。

内容的提问来源于stack exchange，提问作者Kewitschka