Repository: apache/kafka
File: core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala (Scala)

Code before:
      TestUtils.waitUntilTrue(() => (curExpiresPerSec + 1) == DelayedRemoteFetchMetrics.expiredRequestMeter.count(), "DelayedRemoteFetchMetrics.expiredRequestMeter.count() should be 1, but got: " + DelayedRemoteFetchMetrics.expiredRequestMeter.count(), 10000L)
      latch.countDown()
    } finally {
      Utils.tryAll(util.Arrays.asList[Callable[Void]](
        () => {
          replicaManager.shutdown(checkpointHW = false)
          null
        },
        () => {
          remoteLogManager.close()
          null
        }
      ))
    }
  }

  @Test
  def testMultipleRemoteFetchesInOneFetchRequest(): Unit = {
    val replicaId = -1
    val tp0 = new TopicPartition(topic, 0)
    val tp1 = new TopicPartition(topic, 1)
    val tidp0 = new TopicIdPartition(topicId, tp0)
    val tidp1 = new TopicIdPartition(topicId, tp1)

    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), aliveBrokerIds = Seq(0, 1, 2), enableRemoteStorage = true, shouldMockLog = true, remoteFetchQuotaExceeded = Some(false))
    when(replicaManager.logManager.getOrCreateLog(tp0, isNew = false, isFuture = false, Optional.empty(), None).topicId()).thenReturn(Optional.of(topicId))

    try {
      val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints.asJava)
      replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, None)
      replicaManager.createPartition(tp1).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, None)

      val leaderEpoch = 0
      val leaderDelta0 = createLeaderDelta(topicId, tp0, leaderId = 0, leaderEpoch = leaderEpoch)
      val leaderDelta1 = createLeaderDelta(topicId, tp1, leaderId = 0, leaderEpoch = leaderEpoch)
      val leaderMetadataImage0 = imageFromTopics(leaderDelta0.apply())
      val leaderMetadataImage1 = imageFromTopics(leaderDelta1.apply())
      replicaManager.applyDelta(leaderDelta0, leaderMetadataImage0)
      replicaManager.applyDelta(leaderDelta1, leaderMetadataImage1)

      val params = new FetchParams(replicaId, 1, 1000, 10, 100, FetchIsolation.LOG_END, Optional.empty)
      val fetchOffsetTp0 = 1
      val fetchOffsetTp1 = 2

      val responseSeq = new AtomicReference[Seq[(TopicIdPartition, FetchPartitionData)]]()
      val responseLatch = new CountDownLatch(1)

      def fetchCallback(responseStatus: Seq[(TopicIdPartition, FetchPartitionData)]): Unit = {
        responseSeq.set(responseStatus)
        responseLatch.countDown()
      }

Review comment (suggestion):
It is a bit weird to me that `setupMockLog` allows `UnifiedLog` to return "different" topic ids. Perhaps we should enhance `createLogManager` to accept multiple `UnifiedLog`s for the test `testOffsetOutOfRangeExceptionWhenReadFromLog`

Code after:
    } finally {
      Utils.tryAll(util.Arrays.asList[Callable[Void]](
        () => {
          replicaManager.shutdown(checkpointHW = false)
          null
        },
        () => {
          remoteLogManager.close()
          null
        }
      ))
    }
  }

  @Test
  def testMultipleRemoteFetchesInOneFetchRequest(): Unit = {
    val replicaId = -1
    val tp0 = new TopicPartition(topic, 0)
    val tp1 = new TopicPartition(topic, 1)
    val tidp0 = new TopicIdPartition(topicId, tp0)
    val tidp1 = new TopicIdPartition(topicId, tp1)

    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), aliveBrokerIds = Seq(0, 1, 2), enableRemoteStorage = true, shouldMockLog = true, remoteFetchQuotaExceeded = Some(false))

    try {
      val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints.asJava)
      replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, Some(topicId))
      replicaManager.createPartition(tp1).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, Some(topicId))

      val leaderEpoch = 0
      val leaderDelta0 = createLeaderDelta(topicId, tp0, leaderId = 0, leaderEpoch = leaderEpoch)
      val leaderDelta1 = createLeaderDelta(topicId, tp1, leaderId = 0, leaderEpoch = leaderEpoch)
      val leaderMetadataImage0 = imageFromTopics(leaderDelta0.apply())
      val leaderMetadataImage1 = imageFromTopics(leaderDelta1.apply())
      replicaManager.applyDelta(leaderDelta0, leaderMetadataImage0)
      replicaManager.applyDelta(leaderDelta1, leaderMetadataImage1)

      val params = new FetchParams(replicaId, 1, 1000, 10, 100, FetchIsolation.LOG_END, Optional.empty)
      val fetchOffsetTp0 = 1
      val fetchOffsetTp1 = 2

      val responseSeq = new AtomicReference[Seq[(TopicIdPartition, FetchPartitionData)]]()
      val responseLatch = new CountDownLatch(1)

      def fetchCallback(responseStatus: Seq[(TopicIdPartition, FetchPartitionData)]): Unit = {
        responseSeq.set(responseStatus)
        responseLatch.countDown()
      }

      val callbacks: util.Set[Consumer[RemoteLogReadResult]] = new util.HashSet[Consumer[RemoteLogReadResult]]()
      when(mockRemoteLogManager.asyncRead(any(), any())).thenAnswer(ans => {