This is an automated email from the ASF dual-hosted git repository.
pjfanning pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/pekko-connectors-kafka.git
The following commit(s) were added to refs/heads/main by this push:
new ec192ac3 replace waitForDraining (#623)
ec192ac3 is described below
commit ec192ac3a13e41a91a93782141e53b137a40a58c
Author: PJ Fanning <[email protected]>
AuthorDate: Mon Aug 10 09:31:28 2026 +0100
replace waitForDraining (#623)
---
.../internal/TransactionalProducerStage.scala | 30 ++++++-
.../kafka/internal/TransactionalSources.scala | 86 ++++++++++----------
.../apache/pekko/kafka/internal/ProducerSpec.scala | 91 +++++++++++++++++++++-
3 files changed, 157 insertions(+), 50 deletions(-)
diff --git
a/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalProducerStage.scala
b/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalProducerStage.scala
index d41752f7..d34c093d 100644
---
a/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalProducerStage.scala
+++
b/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalProducerStage.scala
@@ -31,6 +31,7 @@ import org.apache.kafka.common.TopicPartition
import scala.concurrent.Future
import scala.concurrent.duration._
import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
/**
* INTERNAL API
@@ -231,12 +232,37 @@ private final class TransactionalProducerStageLogic[K, V,
P](
override def onCompletionSuccess(): Unit = {
log.debug("Committing final transaction before shutdown")
cancelTimer(commitSchedulerKey)
- maybeCommitTransaction(beginNewTransaction = false,
abortEmptyTransactionOnComplete = true)
+ setKeepGoing(true)
+ try {
+ batchOffsets match {
+ case batch: NonemptyTransactionBatch =>
+ commitTransaction(batch, beginNewTransaction = false)
+ case _: EmptyTransactionBatch =>
+ abortTransaction("Transaction is empty and stage is completing")
+ case _ =>
+ ()
+ }
+ } catch {
+ case NonFatal(ex) =>
+ log.error(ex, "Failed to commit final transaction, aborting")
+ try {
+ abortTransaction("Final transaction commit failed")
+ } catch {
+ case NonFatal(abortEx) =>
+ log.error(abortEx, "Failed to abort transaction after commit
failure")
+ }
+ batchOffsets.committingFailed()
+ }
super.onCompletionSuccess()
}
override def onCompletionFailure(ex: Throwable): Unit = {
- abortTransaction("Stage failure")
+ try {
+ abortTransaction("Stage failure")
+ } catch {
+ case NonFatal(abortEx) =>
+ log.error(abortEx, "Failed to abort transaction during stage failure")
+ }
batchOffsets.committingFailed()
super.onCompletionFailure(ex)
}
diff --git
a/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalSources.scala
b/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalSources.scala
index c8ee5b43..907f88d9 100644
---
a/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalSources.scala
+++
b/core/src/main/scala/org/apache/pekko/kafka/internal/TransactionalSources.scala
@@ -36,7 +36,7 @@ import org.apache.kafka.clients.consumer.{ ConsumerConfig,
ConsumerGroupMetadata
import org.apache.kafka.common.{ IsolationLevel, TopicPartition }
import scala.concurrent.duration.FiniteDuration
-import scala.concurrent.{ Await, ExecutionContext, Future }
+import scala.concurrent.{ ExecutionContext, Future }
/** Internal API */
@InternalApi
@@ -168,36 +168,31 @@ private[internal] abstract class
TransactionalSourceLogic[K, V, Msg](shape: Sour
override protected def addToPartitionAssignmentHandler(
handler: PartitionAssignmentHandler): PartitionAssignmentHandler = {
- val blockingRevokedCall = new PartitionAssignmentHandler {
+ val asyncRevokedCall = new PartitionAssignmentHandler {
override def onAssign(assignedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = ()
// This is invoked in the KafkaConsumerActor thread when doing poll.
- override def onRevoke(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
- if (waitForDraining(revokedTps)) {
- sourceActor.ref.tell(Revoked(revokedTps.toList), consumerActor)
- } else {
- sourceActor.ref.tell(Failure(new Error("Timeout while draining")),
consumerActor)
- consumerActor.tell(KafkaConsumerActor.Internal.StopFromStage(id),
consumerActor)
- }
+ // Uses async ask to avoid blocking the poll thread, which would prevent
heartbeats
+ // and could deadlock if the stage actor needs the consumer actor to
process commits.
+ override def onRevoke(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = {
+ import pekko.pattern.ask
+ implicit val timeout: Timeout = Timeout(consumerSettings.commitTimeout)
+ ask(stageActor.ref, Drain(revokedTps, None, Drained))
+ .onComplete {
+ case scala.util.Success(_) =>
+ sourceActor.ref.tell(Revoked(revokedTps.toList), consumerActor)
+ case scala.util.Failure(_) =>
+ sourceActor.ref.tell(Failure(new Error("Timeout while
draining")), consumerActor)
+
consumerActor.tell(KafkaConsumerActor.Internal.StopFromStage(id), consumerActor)
+ }(ExecutionContext.parasitic)
+ }
override def onLost(lostTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
onRevoke(lostTps, consumer)
override def onStop(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = ()
}
- new PartitionAssignmentHelpers.Chain(handler, blockingRevokedCall)
- }
-
- private def waitForDraining(partitions: Set[TopicPartition]): Boolean = {
- import pekko.pattern.ask
- implicit val timeout: Timeout = Timeout(consumerSettings.commitTimeout)
- try {
- Await.result(ask(stageActor.ref, Drain(partitions, None, Drained)),
timeout.duration)
- true
- } catch {
- case t: Throwable =>
- false
- }
+ new PartitionAssignmentHelpers.Chain(handler, asyncRevokedCall)
}
}
@@ -247,19 +242,31 @@ private[kafka] final class TransactionalSubSource[K, V](
override protected def addToPartitionAssignmentHandler(
handler: PartitionAssignmentHandler): PartitionAssignmentHandler = {
- val blockingRevokedCall = new PartitionAssignmentHandler {
+ val asyncRevokedCall = new PartitionAssignmentHandler {
override def onAssign(assignedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = ()
// This is invoked in the KafkaConsumerActor thread when doing poll.
+ // Uses async ask to avoid blocking the poll thread, which would
prevent heartbeats
+ // and could deadlock if the stage actor needs the consumer actor to
process commits.
override def onRevoke(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
- if (revokedTps.isEmpty) ()
- else if (waitForDraining(revokedTps)) {
- subSources.values
- .map(_.controlAndStageActor.stageActor)
- .foreach(_.tell(Revoked(revokedTps.toList), stageActor.ref))
- } else {
- sourceActor.ref.tell(Status.Failure(new Error("Timeout while
draining")), stageActor.ref)
-
consumerActor.tell(KafkaConsumerActor.Internal.StopFromStage(id),
stageActor.ref)
+ if (revokedTps.nonEmpty) {
+ import pekko.pattern.ask
+ implicit val timeout: Timeout =
Timeout(txConsumerSettings.commitTimeout)
+ implicit val ec: ExecutionContext = executionContext
+ Future
+ .sequence(
+ subSources.values.map(_.stageActor).map(ask(_,
Drain(revokedTps, None, Drained))))
+ .onComplete {
+ case scala.util.Success(_) =>
+ subSources.values
+ .map(_.controlAndStageActor.stageActor)
+ .foreach(_.tell(Revoked(revokedTps.toList),
stageActor.ref))
+ case scala.util.Failure(_) =>
+ sourceActor.ref.tell(
+ Status.Failure(new Error("Timeout while draining")),
+ stageActor.ref)
+
consumerActor.tell(KafkaConsumerActor.Internal.StopFromStage(id),
stageActor.ref)
+ }
}
override def onLost(lostTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
@@ -267,22 +274,7 @@ private[kafka] final class TransactionalSubSource[K, V](
override def onStop(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = ()
}
- new PartitionAssignmentHelpers.Chain(handler, blockingRevokedCall)
- }
-
- private def waitForDraining(partitions: Set[TopicPartition]): Boolean = {
- import pekko.pattern.ask
- implicit val timeout: Timeout =
Timeout(txConsumerSettings.commitTimeout)
- try {
- val drainCommandFutures =
- subSources.values.map(_.stageActor).map(ask(_, Drain(partitions,
None, Drained)))
- implicit val ec: ExecutionContext = executionContext
- Await.result(Future.sequence(drainCommandFutures), timeout.duration)
- true
- } catch {
- case t: Throwable =>
- false
- }
+ new PartitionAssignmentHelpers.Chain(handler, asyncRevokedCall)
}
}
}
diff --git
a/tests/src/test/scala/org/apache/pekko/kafka/internal/ProducerSpec.scala
b/tests/src/test/scala/org/apache/pekko/kafka/internal/ProducerSpec.scala
index 53cbe2d1..e064baf0 100644
--- a/tests/src/test/scala/org/apache/pekko/kafka/internal/ProducerSpec.scala
+++ b/tests/src/test/scala/org/apache/pekko/kafka/internal/ProducerSpec.scala
@@ -30,7 +30,7 @@ import pekko.{ Done, NotUsed }
import com.typesafe.config.ConfigFactory
import org.apache.kafka.clients.consumer.{ ConsumerGroupMetadata,
OffsetAndMetadata }
import org.apache.kafka.clients.producer._
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.{ KafkaException, TopicPartition }
import org.apache.kafka.common.serialization.StringSerializer
import org.mockito
import org.mockito.Mockito
@@ -563,6 +563,80 @@ class ProducerSpec(_system: ActorSystem)
client.verifyTxAbort()
client.verifyClosed()
}
+
+ it should "abort transaction and close producer when commitTransaction fails
during shutdown" in {
+ val input = recordAndMetadata(1)
+
+ val client = {
+ val inputMap = Map(input)
+ new ProducerMock[K, V](ProducerMock.handlers.delayedMap(100.millis)(x =>
Try { inputMap(x) }))
+ }
+ val committedMarker = new CommittedMarkerMock
+
+ val (source, sink) = TestSource[TxMsg]()
+ .via(testTransactionProducerFlow(client))
+ .toMat(TestSink())(Keep.both)
+ .run()
+
+ val txMsg: TxMsg = toTxMessage(input, committedMarker.mock)
+ source.sendNext(txMsg)
+ sink.requestNext()
+
+ client.verifySend(atLeastOnce())
+
+ // Wait for the timer-based commit to succeed first
+ awaitAssert(client.verifyTxCommit(txMsg.passThrough), 2.second)
+
+ // Now make commitTransaction throw for the shutdown commit
+ Mockito
+ .doAnswer(_ => throw new KafkaException("commit failed"))
+ .when(client.mock)
+ .commitTransaction()
+
+ // Complete the source to trigger onCompletionSuccess
+ source.sendComplete()
+ sink.expectComplete()
+
+ // Verify: commitTransaction was called (and threw), then abortTransaction
as fallback, then producer closed
+ client.verifyTxAbortAfterFailedCommit()
+ }
+
+ it should "close producer when abortTransaction fails during stage failure"
in {
+ val input = recordAndMetadata(1)
+
+ val client = {
+ val inputMap = Map(input)
+ new ProducerMock[K, V](ProducerMock.handlers.delayedMap(100.millis)(x =>
Try { inputMap(x) }))
+ }
+ val committedMarker = new CommittedMarkerMock
+
+ val (source, sink) = TestSource[TxMsg]()
+ .via(testTransactionProducerFlow(client))
+ .toMat(Sink.lastOption)(Keep.both)
+ .run()
+
+ val txMsg = toTxMessage(input, committedMarker.mock)
+ source.sendNext(txMsg)
+
+ awaitAssert(client.verifyTxInitialized())
+
+ // Make abortTransaction throw
+ Mockito
+ .doAnswer(_ => throw new KafkaException("abort failed"))
+ .when(client.mock)
+ .abortTransaction()
+
+ // Trigger stage failure
+ source.sendError(new Exception("upstream failure"))
+
+ Await.ready(sink, remainingOrDefault)
+ sink.value should matchPattern {
+ case Some(Failure(_)) =>
+ }
+
+ // Even though abortTransaction throws, the producer should still be closed
+ client.verifyClosedAfterFailedAbort()
+ }
}
object ProducerMock {
@@ -672,6 +746,21 @@ class ProducerMock[K, V](handler: ProducerMock.Handler[K,
V])(implicit ec: Execu
inOrder.verify(mock).flush()
inOrder.verify(mock).close(mockito.ArgumentMatchers.any[java.time.Duration])
}
+
+ def verifyTxAbortAfterFailedCommit() = {
+ val inOrder = Mockito.inOrder(mock)
+ inOrder.verify(mock).commitTransaction()
+ inOrder.verify(mock).abortTransaction()
+ inOrder.verify(mock).flush()
+
inOrder.verify(mock).close(mockito.ArgumentMatchers.any[java.time.Duration])
+ }
+
+ def verifyClosedAfterFailedAbort() = {
+ val inOrder = Mockito.inOrder(mock)
+ inOrder.verify(mock).abortTransaction()
+ inOrder.verify(mock).flush()
+
inOrder.verify(mock).close(mockito.ArgumentMatchers.any[java.time.Duration])
+ }
}
class CommittedMarkerMock {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
