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 1 commit
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
99 changes: 50 additions & 49 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,68 @@ 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 =
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
}
)
@inline def insertSentResult(resultIndex: Int, sentResult: Either[Throwable, RecordMetadata]): Unit = {
var notInserted = true

while (notInserted) {
val sentRecordsCount = sentRecordsCounter.get()

// Updating sentResults[resultIndex] here is safe,
// because volatile variable read performed right before the update and volatile variable write after,
// 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 after volatile variable read
sentResults.update(resultIndex, sentResult)
Copy link
Collaborator

@erikvanoosten erikvanoosten 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.

Since sentResults is an Array, we can assign to any index, as long as we know there is a single writer for each index (which is the case). The only thing that needs to be multi-thread protected is sentRecordsCounter but that is now an AtomicInteger, so calling increaseAndGet would be sufficient, no while loop needed.

Or did I miss something?

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.

Yeah, you are right Erik, thats awesome 🔥
Many thanks for the suggestion!
I was afraid that without having the latest updated version of sentResults during current update, we may run into issue with losing concurrent update result to some other index (similar to word tearing), but from some googling I can see that JVM handles this, and it is safe to do as you suggesting.


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

Unsafe.unsafe { implicit u =>
exec {
runtime.unsafe.run(done.succeed(sentResultsChunk))
}
}
}

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
notInserted = !sentRecordsCounter.compareAndSet(sentRecordsCount, newSentRecordsCount)
}
}

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 Down