Repository: apache/kafka
File: core/src/main/scala/kafka/server/ReplicaManager.scala (Scala)

Code before:
  private def processRemoteFetches(remoteFetchInfos: util.HashMap[TopicIdPartition, RemoteStorageFetchInfo],
                                   params: FetchParams,
                                   responseCallback: Seq[(TopicIdPartition, FetchPartitionData)] => Unit,
                                   logReadResults: Seq[(TopicIdPartition, LogReadResult)],
                                   remoteFetchPartitionStatus: Seq[(TopicIdPartition, FetchPartitionStatus)]): Unit = {
    val remoteFetchTasks = new util.HashMap[TopicIdPartition, Future[Void]]
    val remoteFetchResults = new util.HashMap[TopicIdPartition, CompletableFuture[RemoteLogReadResult]]

    remoteFetchInfos.forEach { (topicIdPartition, remoteFetchInfo) =>
      val (task, result) = processRemoteFetch(remoteFetchInfo)
      remoteFetchTasks.put(topicIdPartition, task)
      remoteFetchResults.put(topicIdPartition, result)
    }

    val remoteFetchMaxWaitMs = config.remoteLogManagerConfig.remoteFetchMaxWaitMs().toLong
    val remoteFetch = new DelayedRemoteFetch(remoteFetchTasks, remoteFetchResults, remoteFetchInfos, remoteFetchMaxWaitMs,
      remoteFetchPartitionStatus, params, logReadResults, this, responseCallback)

    // create a list of (topic, partition) pairs to use as keys for this delayed fetch operation
    val delayedFetchKeys = remoteFetchTasks.asScala.map { case (tp, _) => new TopicPartitionOperationKey(tp) }.toList
    delayedRemoteFetchPurgatory.tryCompleteElseWatch(remoteFetch, delayedFetchKeys.asJava)
  }

  /**
   * Fetch messages from a replica, and wait until enough data can be fetched and return;
   * the callback function will be triggered either when timeout or required fetch info is satisfied.
   * Consumers may fetch from any replica, but followers can only fetch from the leader.
   */
  def fetchMessages(params: FetchParams,
                    fetchInfos: Seq[(TopicIdPartition, PartitionData)],
                    quota: ReplicaQuota,
                    responseCallback: Seq[(TopicIdPartition, FetchPartitionData)] => Unit): Unit = {

    // check if this fetch request can be satisfied right away
    val logReadResults = readFromLog(params, fetchInfos, quota, readFromPurgatory = false)
    var bytesReadable: Long = 0
    var errorReadingData = false

    // topic-partitions that have to be read from remote storage
    val remoteFetchInfos = new util.HashMap[TopicIdPartition, RemoteStorageFetchInfo]()

    var hasDivergingEpoch = false
    var hasPreferredReadReplica = false
    val logReadResultMap = new mutable.HashMap[TopicIdPartition, LogReadResult]

    logReadResults.foreach { case (topicIdPartition, logReadResult) =>
      brokerTopicStats.topicStats(topicIdPartition.topicPartition.topic).totalFetchRequestRate.mark()
      brokerTopicStats.allTopicsStats.totalFetchRequestRate.mark()
      if (logReadResult.error != Errors.NONE)
        errorReadingData = true
      if (logReadResult.info.delayedRemoteStorageFetch.isPresent) {

Review comment (performance):
Sorry, I don't get the fix. I'm really surprised this fixes the memory leak. Could you explain why it leaks memory here? I thought we only use the `tp` in `remoteFetchPartitionStatus` to create the delayedFetchKeys list, where is the leak? We don't hold the `FetchPartitionStatus` at all, right? 

Code after:
    (remoteFetchTask, remoteFetchResult)
  }

  /**
   * Process all remote fetches by creating async read tasks and handling them in DelayedRemoteFetch collectively.
   */
  private def processRemoteFetches(remoteFetchInfos: util.HashMap[TopicIdPartition, RemoteStorageFetchInfo],
                                   params: FetchParams,
                                   responseCallback: Seq[(TopicIdPartition, FetchPartitionData)] => Unit,
                                   logReadResults: Seq[(TopicIdPartition, LogReadResult)],
                                   fetchPartitionStatus: Seq[(TopicIdPartition, FetchPartitionStatus)]): Unit = {
    val remoteFetchTasks = new util.HashMap[TopicIdPartition, Future[Void]]
    val remoteFetchResults = new util.HashMap[TopicIdPartition, CompletableFuture[RemoteLogReadResult]]

    remoteFetchInfos.forEach { (topicIdPartition, remoteFetchInfo) =>
      val (task, result) = processRemoteFetch(remoteFetchInfo)
      remoteFetchTasks.put(topicIdPartition, task)
      remoteFetchResults.put(topicIdPartition, result)
    }

    val remoteFetchMaxWaitMs = config.remoteLogManagerConfig.remoteFetchMaxWaitMs().toLong
    val remoteFetch = new DelayedRemoteFetch(remoteFetchTasks, remoteFetchResults, remoteFetchInfos, remoteFetchMaxWaitMs,
      fetchPartitionStatus, params, logReadResults, this, responseCallback)

    // create a list of (topic, partition) pairs to use as keys for this delayed fetch operation
    val delayedFetchKeys = remoteFetchTasks.asScala.map { case (tp, _) => new TopicPartitionOperationKey(tp) }.toList
    delayedRemoteFetchPurgatory.tryCompleteElseWatch(remoteFetch, delayedFetchKeys.asJava)
  }

  /**
   * Fetch messages from a replica, and wait until enough data can be fetched and return;
   * the callback function will be triggered either when timeout or required fetch info is satisfied.
   * Consumers may fetch from any replica, but followers can only fetch from the leader.
   */
  def fetchMessages(params: FetchParams,
                    fetchInfos: Seq[(TopicIdPartition, PartitionData)],
                    quota: ReplicaQuota,
                    responseCallback: Seq[(TopicIdPartition, FetchPartitionData)] => Unit): Unit = {

    // check if this fetch request can be satisfied right away
    val logReadResults = readFromLog(params, fetchInfos, quota, readFromPurgatory = false)
    var bytesReadable: Long = 0
    var errorReadingData = false

    // topic-partitions that have to be read from remote storage
    val remoteFetchInfos = new util.HashMap[TopicIdPartition, RemoteStorageFetchInfo]()

    var hasDivergingEpoch = false
    var hasPreferredReadReplica = false
    val logReadResultMap = new mutable.HashMap[TopicIdPartition, LogReadResult]