Repository: apache/kafka
File: core/src/test/scala/integration/kafka/coordinator/transaction/ProducerIntegrationTest.scala (Scala)

Code before:
  private def verifyUniqueIds(clusterInstance: ClusterInstance): Unit = {
    // Request enough PIDs from each broker to ensure each broker generates two blocks
    val ids = clusterInstance.brokerSocketServers().stream().flatMap( broker => {
      IntStream.range(0, 1001).parallel().mapToObj( _ =>
        nextProducerId(broker, clusterInstance.clientListener())
      )}).collect(Collectors.toList[Long]).asScala.toSeq

    val brokerCount = clusterInstance.brokerIds.size
    val expectedTotalCount = 1001 * brokerCount
    assertEquals(expectedTotalCount, ids.size, s"Expected exactly $expectedTotalCount IDs")
    assertEquals(expectedTotalCount, ids.distinct.size, "Found duplicate producer IDs")
  }

  private def nextProducerId(broker: SocketServer, listener: ListenerName): Long = {
    // Generating producer ids may fail while waiting for the initial block and also
    // when the current block is full and waiting for the prefetched block.
    val deadline = 5.seconds.fromNow
    var shouldRetry = true
    var response: InitProducerIdResponse = null
    while (shouldRetry && deadline.hasTimeLeft()) {
      val data = new InitProducerIdRequestData()
        .setProducerEpoch(RecordBatch.NO_PRODUCER_EPOCH)
        .setProducerId(RecordBatch.NO_PRODUCER_ID)
        .setTransactionalId(null)
        .setTransactionTimeoutMs(10)
    val request = new InitProducerIdRequest.Builder(data).build()
    val port = broker.boundPort(listener)

      response = IntegrationTestUtils.connectAndReceive[InitProducerIdResponse](request,
        port
        )

      shouldRetry = response.data.errorCode == Errors.COORDINATOR_LOAD_IN_PROGRESS.code
    }
    assertTrue(deadline.hasTimeLeft())
    assertEquals(Errors.NONE.code, response.data.errorCode)
    response.data().producerId()
  }
}


Review comment (suggestion):
```scala
      val request = new InitProducerIdRequest.Builder(data).build()
      val port = broker.boundPort(listener)
      response = IntegrationTestUtils.connectAndReceive[InitProducerIdResponse](request, port)
      shouldRetry = response.data.errorCode == Errors.COORDINATOR_LOAD_IN_PROGRESS.code
```

Code after:
  private def verifyUniqueIds(clusterInstance: ClusterInstance): Unit = {
    // Request enough PIDs from each broker to ensure each broker generates two blocks
    val ids = clusterInstance.brokerSocketServers().stream().flatMap( broker => {
      IntStream.range(0, 1001).parallel().mapToObj( _ =>
        nextProducerId(broker, clusterInstance.clientListener())
      )}).collect(Collectors.toList[Long]).asScala.toSeq

    val brokerCount = clusterInstance.brokerIds.size
    val expectedTotalCount = 1001 * brokerCount
    assertEquals(expectedTotalCount, ids.size, s"Expected exactly $expectedTotalCount IDs")
    assertEquals(expectedTotalCount, ids.distinct.size, "Found duplicate producer IDs")
  }

  private def nextProducerId(broker: SocketServer, listener: ListenerName): Long = {
    // Generating producer ids may fail while waiting for the initial block and also
    // when the current block is full and waiting for the prefetched block.
    val deadline = 5.seconds.fromNow
    var shouldRetry = true
    var response: InitProducerIdResponse = null
    while (shouldRetry && deadline.hasTimeLeft()) {
      val data = new InitProducerIdRequestData()
        .setProducerEpoch(RecordBatch.NO_PRODUCER_EPOCH)
        .setProducerId(RecordBatch.NO_PRODUCER_ID)
        .setTransactionalId(null)
        .setTransactionTimeoutMs(10)
      val request = new InitProducerIdRequest.Builder(data).build()
      val port = broker.boundPort(listener)
      response = IntegrationTestUtils.connectAndReceive[InitProducerIdResponse](request, port)
      shouldRetry = response.data.errorCode == Errors.COORDINATOR_LOAD_IN_PROGRESS.code
    }
    assertTrue(deadline.hasTimeLeft())
    assertEquals(Errors.NONE.code, response.data.errorCode)
    response.data().producerId()
  }
}