Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize Producer sendFromQueue implementation #1326

Merged
merged 2 commits into from
Sep 17, 2024
Merged

Optimize Producer sendFromQueue implementation #1326

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0e94d7a
Optimize Producer sendFromQueue implementation
ytalashko Sep 14, 2024
107405a
Remove redudnat loop from produce callback in Producer.sendFromQueue
ytalashko Sep 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 38 additions & 52 deletions zio-kafka/src/main/scala/zio/kafka/producer/Producer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import zio.kafka.serde.Serializer
import zio.kafka.utils.SslHelper
import zio.stream.{ ZPipeline, ZStream }

import java.util.concurrent.atomic.AtomicInteger
import scala.jdk.CollectionConverters._
import scala.util.control.{ NoStackTrace, NonFatal }

Expand Down Expand Up @@ -484,68 +485,56 @@ private[producer] final class ProducerLive(
ZStream
.fromQueueWithShutdown(sendQueue)
.mapZIO { case (serializedRecords, done) =>
ZIO.suspendSucceed {
ZIO.succeed {
val recordsLength = serializedRecords.length
val sentRecordsCounter = new AtomicInteger(0)
val recordsIterator: Iterator[(ByteRecord, Int)] = serializedRecords.iterator.zipWithIndex
val sentResults: Array[Either[Throwable, RecordMetadata]] =
new Array[Either[Throwable, RecordMetadata]](recordsLength)

Ref.make(0).map { sentRecordsCountRef =>
@inline def safelyInsertSentResult(resultIndex: Int, sentResult: Either[Throwable, RecordMetadata]): Unit =
@inline def insertSentResult(resultIndex: Int, sentResult: Either[Throwable, RecordMetadata]): Unit = {
// Updating sentResults[resultIndex] here is safe,
// because only a single update for every resultIndex of sentResults is performed
sentResults.update(resultIndex, sentResult)

// Reading from sentRecordsCounter guarantees fully updated version of sentResults
// will be visible and used to complete done promise
if (sentRecordsCounter.incrementAndGet() == recordsLength) {
val sentResultsChunk = Chunk.fromArray(sentResults)

Unsafe.unsafe { implicit u =>
exec {
runtime.unsafe.run(
sentRecordsCountRef.update { sentRecordsCount =>
// Updating sentResults[resultIndex] here is safe,
// cause Ref.update starts with volatile variable read and ends with volatile variable write,
// which guarantees sentResults.update executed on the latest updated version of sentResults
// and currently updated version of sentResults
// will be visible to the next sentResults read or update called within Ref.update
sentResults.update(resultIndex, sentResult)

val newSentRecordsCount = sentRecordsCount + 1
if (newSentRecordsCount == recordsLength) {
val sentResultsChunk = Chunk.fromArray(sentResults)

exec {
runtime.unsafe.run(done.succeed(sentResultsChunk))
}
}

newSentRecordsCount
}
)
}
val _ = runtime.unsafe.run(done.succeed(sentResultsChunk))
Copy link
Contributor Author

@ytalashko ytalashko Sep 15, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inlined exec here, as it was the only place of usage

}
}
}

var previousSendCallsSucceed = true

recordsIterator.foreach { case (record: ByteRecord, recordIndex: Int) =>
if (previousSendCallsSucceed) {
try {
val _ = p.send(
record,
(metadata: RecordMetadata, err: Exception) =>
safelyInsertSentResult(
recordIndex,
if (err eq null) Right(metadata) else Left(err)
)
)
} catch {
case NonFatal(err) =>
previousSendCallsSucceed = false
var previousSendCallsSucceed = true

safelyInsertSentResult(
recordsIterator.foreach { case (record: ByteRecord, recordIndex: Int) =>
if (previousSendCallsSucceed) {
try {
val _ = p.send(
record,
(metadata: RecordMetadata, err: Exception) =>
insertSentResult(
recordIndex,
Left(err)
if (err eq null) Right(metadata) else Left(err)
)
}
} else {
safelyInsertSentResult(
recordIndex,
Left(Producer.PublishOmittedException)
)
} catch {
case NonFatal(err) =>
previousSendCallsSucceed = false

insertSentResult(
recordIndex,
Left(err)
)
}
} else {
insertSentResult(
recordIndex,
Left(Producer.PublishOmittedException)
)
}
}
}
Expand All @@ -561,7 +550,4 @@ private[producer] final class ProducerLive(
key <- keySerializer.serialize(r.topic, r.headers, r.key())
value <- valueSerializer.serialize(r.topic, r.headers, r.value())
} yield new ProducerRecord(r.topic, r.partition(), r.timestamp(), key, value, r.headers)

/** Used to prevent warnings about not using the result of an expression. */
@inline private def exec[A](f: => A): Unit = { val _ = f }
}