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

Code before:
      }, "poll records size not match")
      val lastRecord = StreamSupport.stream(records.spliterator, false).reduce((_, second) => second).orElse(null)
      val offsets = util.Map.of(
        new TopicPartition(lastRecord.topic, lastRecord.partition), new OffsetAndMetadata(lastRecord.offset + 1))
      producer.sendOffsetsToTransaction(offsets, consumer.groupMetadata)
      producer.commitTransaction()
    } finally {
      if (producer != null) producer.close()
      if (consumer != null) consumer.close()
    }

    val admin: Admin = cluster.admin
    try {
      TestUtils.waitForCondition(() => {
        admin.listTransactions.all.get.stream
          .filter(txn => txn.transactionalId == txnId)
          .anyMatch(txn => txn.state eq TransactionState.COMPLETE_COMMIT)
      }, "transaction is not in COMPLETE_COMMIT state")
    } finally if (admin != null) admin.close()
  }

  private def verifyUniqueIds(clusterInstance: ClusterInstance): Unit = {
    // Request enough PIDs from each broker to ensure each broker generates two blocks
    val ids = clusterInstance.brokers().values().stream().flatMap( broker => {
      IntStream.range(0, 1001).parallel().mapToObj( _ =>
        nextProducerId(broker.socketServer, 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
    }

Review comment (refactor):
Could you pleas change the arguments of `nextProducerId` to "port"? that can simplify the method and then we don't need to use `SocketServer`

Code after:
      val offsets = util.Map.of(
        new TopicPartition(lastRecord.topic, lastRecord.partition), new OffsetAndMetadata(lastRecord.offset + 1))
      producer.sendOffsetsToTransaction(offsets, consumer.groupMetadata)
      producer.commitTransaction()
    } finally {
      if (producer != null) producer.close()
      if (consumer != null) consumer.close()
    }

    val admin: Admin = cluster.admin
    try {
      TestUtils.waitForCondition(() => {
        admin.listTransactions.all.get.stream
          .filter(txn => txn.transactionalId == txnId)
          .anyMatch(txn => txn.state eq TransactionState.COMPLETE_COMMIT)
      }, "transaction is not in COMPLETE_COMMIT state")
    } finally if (admin != null) admin.close()
  }

  private def verifyUniqueIds(clusterInstance: ClusterInstance): Unit = {
    // Request enough PIDs from each broker to ensure each broker generates two blocks
    val ids = clusterInstance.brokers().values().stream().flatMap(broker => {
      IntStream.range(0, 1001).parallel().mapToObj(_ =>
        nextProducerId(broker.boundPort(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(port: Int): 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()
      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()