Optionally commit offsets synchronously on partition revoke

build.sbt

@@ -270,6 +270,11 @@ ThisBuild / mimaBinaryIssueFilters ++= {
     ProblemFilters.exclude[ReversedMissingMethodProblem]("fs2.kafka.KafkaConsumer.settings"),
     ProblemFilters.exclude[ReversedMissingMethodProblem]("fs2.kafka.ConsumerSettings.withMaxParallelism"),
     ProblemFilters.exclude[ReversedMissingMethodProblem]("fs2.kafka.ConsumerSettings.maxParallelism"),
+    // `commitOnRevoke` settings flag, added (opt-in) in 4.1. `ConsumerSettings` is a sealed
+    // abstract class (no external implementors), so adding these abstract methods is safe; the
+    // `ConsumerSettingsImpl` copy/this/apply filters below already cover the extra field.
+    ProblemFilters.exclude[ReversedMissingMethodProblem]("fs2.kafka.ConsumerSettings.commitOnRevoke"),
+    ProblemFilters.exclude[ReversedMissingMethodProblem]("fs2.kafka.ConsumerSettings.withCommitOnRevoke"),
     ProblemFilters.exclude[DirectMissingMethodProblem]("fs2.kafka.ConsumerSettings#ConsumerSettingsImpl.copy"),
     ProblemFilters.exclude[DirectMissingMethodProblem]("fs2.kafka.ConsumerSettings#ConsumerSettingsImpl.this"),
     ProblemFilters.exclude[DirectMissingMethodProblem]("fs2.kafka.ConsumerSettings#ConsumerSettingsImpl.apply"),

modules/core/src/main/scala/fs2/kafka/ConsumerSettings.scala

@@ -414,6 +414,27 @@ sealed abstract class ConsumerSettings[F[_], K, V] {
     */
   def maxParallelism: Int
 
+  /**
+    * Whether, when partitions are revoked, the offsets of already-requested commits for those
+    * partitions should be committed synchronously (via `commitSync`) from within the rebalance
+    * listener, while the consumer still owns them.
+    *
+    * The default is `false`. When `false`, the behaviour is unchanged: commits are issued
+    * asynchronously and an in-flight commit may be lost if a rebalance or shutdown happens before
+    * it is acknowledged, leading to records being re-delivered (at-least-once).
+    *
+    * When `true`, the latest requested offset per partition is remembered and re-committed
+    * synchronously on revoke, which narrows that duplicate-delivery window. It does not eliminate
+    * duplicates entirely (records consumed but not yet committed are still re-delivered), and is
+    * best combined with [[RebalanceRevokeMode.Graceful]] and/or idempotent processing.
+    */
+  def commitOnRevoke: Boolean
+
+  /**
+    * Creates a new [[ConsumerSettings]] with the specified [[commitOnRevoke]] flag.
+    */
+  def withCommitOnRevoke(commitOnRevoke: Boolean): ConsumerSettings[F, K, V]
+
 }
 
 object ConsumerSettings {
@@ -431,7 +452,8 @@ object ConsumerSettings {
     override val recordMetadata: ConsumerRecord[K, V] => String,
     override val maxPrefetchBatches: Int,
     override val rebalanceRevokeMode: RebalanceRevokeMode,
-    override val maxParallelism: Int
+    override val maxParallelism: Int,
+    override val commitOnRevoke: Boolean
   ) extends ConsumerSettings[F, K, V] {
 
     override def withMaxParallelism(maxParallelism: Int): ConsumerSettings[F, K, V] =
@@ -568,6 +590,9 @@ object ConsumerSettings {
     ): ConsumerSettings[F, K, V] =
       copy(rebalanceRevokeMode = rebalanceRevokeMode)
 
+    override def withCommitOnRevoke(commitOnRevoke: Boolean): ConsumerSettings[F, K, V] =
+      copy(commitOnRevoke = commitOnRevoke)
+
     override def toString: String =
       s"ConsumerSettings(closeTimeout = $closeTimeout, commitTimeout = $commitTimeout, pollInterval = $pollInterval, pollTimeout = $pollTimeout, commitRecovery = $commitRecovery)"
 
@@ -603,7 +628,8 @@ object ConsumerSettings {
       recordMetadata = _ => OffsetFetchResponse.NO_METADATA,
       maxPrefetchBatches = 2,
       rebalanceRevokeMode = RebalanceRevokeMode.Eager,
-      maxParallelism = Int.MaxValue
+      maxParallelism = Int.MaxValue,
+      commitOnRevoke = false
     )
 
   def apply[F[_], K, V](

modules/core/src/main/scala/fs2/kafka/internal/LogEntry.scala

@@ -18,6 +18,7 @@ import fs2.kafka.internal.syntax.*
 import fs2.kafka.internal.LogLevel.*
 import fs2.kafka.CommittableConsumerRecord
 
+import org.apache.kafka.clients.consumer.OffsetAndMetadata
 import org.apache.kafka.common.TopicPartition
 
 sealed abstract private[kafka] class LogEntry {
@@ -124,6 +125,26 @@ private[kafka] object LogEntry {
 
   }
 
+  final case class CommittedOffsetsOnRevoke(
+    revoked: Set[TopicPartition],
+    offsets: Map[TopicPartition, OffsetAndMetadata],
+    result: Either[Throwable, Unit]
+  ) extends LogEntry {
+
+    override def level: LogLevel = Info
+
+    override def message: String =
+      result match {
+        case Right(()) =>
+          s"Committed offsets [${offsets
+              .mkString(", ")}] synchronously on revoke of partitions [${revoked.mkString(", ")}]."
+        case Left(e) =>
+          s"Failed to commit offsets [${offsets
+              .mkString(", ")}] synchronously on revoke of partitions [${revoked.mkString(", ")}]: $e."
+      }
+
+  }
+
   def recordsString[F[_]](
     records: Map[Set[TopicPartition], List[CommittableConsumerRecord[F, ?, ?]]]
   ): String =

modules/core/src/main/scala/fs2/kafka/internal/WithConsumer.scala

@@ -15,7 +15,20 @@ import fs2.kafka.internal.syntax.*
 import org.apache.kafka.clients.consumer.CloseOptions
 
 sealed abstract private[kafka] class WithConsumer[F[_]] {
+
   def blocking[A](f: KafkaByteConsumer => A): F[A]
+
+  /**
+    * Runs `f` on the current thread, directly on the underlying consumer, bypassing the
+    * single-threaded blocking context used by [[blocking]].
+    *
+    * This is ONLY safe to call from within a `ConsumerRebalanceListener` callback. Those callbacks
+    * are invoked by Kafka on the consumer's polling thread (inside `poll`), so the consumer may be
+    * accessed reentrantly from there, and routing through [[blocking]] would instead deadlock by
+    * submitting to the single-threaded context that is already occupied by the in-progress `poll`.
+    */
+  def synchronouslyDuringRebalance[A](f: KafkaByteConsumer => A): A
+
 }
 
 private[kafka] object WithConsumer {
@@ -37,6 +50,9 @@ private[kafka] object WithConsumer {
             override def blocking[A](f: KafkaByteConsumer => A): F[A] =
               b(f(consumer))
 
+            override def synchronouslyDuringRebalance[A](f: KafkaByteConsumer => A): A =
+              f(consumer)
+
           }
         }
       }(_.blocking(_.close(CloseOptions.timeout(settings.closeTimeout.toJava))))

modules/core/src/main/scala/fs2/kafka/internal/actor/KafkaConsumerActor.scala

@@ -10,6 +10,7 @@ import java.time.Duration
 import java.util
 
 import scala.collection.immutable.SortedSet
+import scala.util.control.NonFatal
 import scala.util.matching.Regex
 
 import cats.data.NonEmptyList
@@ -91,14 +92,17 @@ final private[kafka] class KafkaConsumerActor[F[_], K, V](
   private[this] val maxPrefetch: Int      = settings.maxPrefetchBatches
 
   val consumerRebalanceListener: ConsumerRebalanceListener = new ConsumerRebalanceListener {
-    override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit =
+    override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+      if (settings.commitOnRevoke)
+        commitOnRevokeSync(partitions.asScala.toSet)
       dispatcher.unsafeRunSync {
         state.evalUpdate { state =>
           val currentAssignment = state.partitionGroupState.keys.toList.flatten.toSet
           val targetAssignment  = currentAssignment -- partitions.asScala
           alignPartitionState(targetAssignment, state.partitionGroupState).map(state.withGroupState)
         }
       }
+    }
 
     override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit =
       dispatcher.unsafeRunSync {
@@ -381,7 +385,39 @@ final private[kafka] class KafkaConsumerActor[F[_], K, V](
       .handleErrorWith(e => F.delay(callback(Left(e))))
 
   private[this] def commit(request: Request.Commit[F]): F[Unit] =
-    commitAsync(request.offsets, request.callback)
+    state
+      .update(_.withRequestedCommitOffsets(request.offsets))
+      .whenA(settings.commitOnRevoke) >> commitAsync(request.offsets, request.callback)
+
+  /**
+    * Best-effort, synchronous commit of the most recently requested offsets for the partitions
+    * being revoked, performed from within `onPartitionsRevoked` while we still own them.
+    *
+    * Kafka invokes the rebalance listener on the consumer's polling thread, from inside `poll`. The
+    * `commitSync` therefore goes through [[WithConsumer.synchronouslyDuringRebalance]] — a direct,
+    * reentrant call on the Java consumer on that same thread — rather than `withConsumer.blocking`,
+    * which would submit to the single-threaded blocking context already occupied by `poll` and
+    * deadlock. State access is bounced onto the effect runtime (it is thread-safe), but the
+    * consumer call itself must stay on the polling thread.
+    *
+    * Failures are logged and swallowed: this only narrows the duplicate-delivery window, and the
+    * consumer remains at-least-once.
+    */
+  private[this] def commitOnRevokeSync(revoked: Set[TopicPartition]): Unit = {
+    val offsets = dispatcher.unsafeRunSync(state.modify(_.removeRequestedCommitOffsets(revoked)))
+    if (offsets.nonEmpty) {
+      val result =
+        try {
+          withConsumer.synchronouslyDuringRebalance(
+            _.commitSync(offsets.asJava, settings.commitTimeout.toJava)
+          )
+          ().asRight[Throwable]
+        } catch {
+          case NonFatal(e) => Left(e)
+        }
+      dispatcher.unsafeRunSync(logging.log(CommittedOffsetsOnRevoke(revoked, offsets, result)))
+    }
+  }
 
   private[this] def manualCommitSync(request: Request.ManualCommitSync[F]): F[Unit] = {
     val commit =

modules/core/src/main/scala/fs2/kafka/internal/actor/State.scala

@@ -13,6 +13,7 @@ import cats.effect.Deferred
 import fs2.kafka.CommittableConsumerRecord
 import fs2.Chunk
 
+import org.apache.kafka.clients.consumer.OffsetAndMetadata
 import org.apache.kafka.common.TopicPartition
 
 private[kafka] object State {
@@ -21,7 +22,8 @@ private[kafka] object State {
     State[F, K, V](
       Map.empty,
       false,
-      false
+      false,
+      Map.empty
     )
 
 }
@@ -46,7 +48,8 @@ final private[kafka] case class PartitionGroupState[F[_], K, V](
 final private[kafka] case class State[F[_], K, V](
   partitionGroupState: Map[Set[TopicPartition], PartitionGroupState[F, K, V]],
   subscribed: Boolean,
-  streaming: Boolean
+  streaming: Boolean,
+  requestedCommitOffsets: Map[TopicPartition, OffsetAndMetadata]
 )(implicit
   F: Async[F]
 ) {
@@ -64,6 +67,31 @@ final private[kafka] case class State[F[_], K, V](
 
   def withNotStreaming(): State[F, K, V] = copy(streaming = false)
 
+  /**
+    * Remembers the latest requested commit offset per partition, so that — if those partitions are
+    * revoked before the asynchronous commit is acknowledged — they can be committed synchronously
+    * from within the rebalance listener. Keeps the highest offset seen per partition.
+    */
+  def withRequestedCommitOffsets(
+    offsets: Map[TopicPartition, OffsetAndMetadata]
+  ): State[F, K, V] =
+    copy(requestedCommitOffsets = offsets.foldLeft(requestedCommitOffsets) {
+      case (acc, (partition, offsetAndMetadata)) =>
+        val isNewer = acc.get(partition).forall(_.offset < offsetAndMetadata.offset)
+        if (isNewer) acc.updated(partition, offsetAndMetadata) else acc
+    })
+
+  /**
+    * Removes and returns the tracked commit offsets for the given (revoked) partitions.
+    */
+  def removeRequestedCommitOffsets(
+    revoked: Set[TopicPartition]
+  ): (State[F, K, V], Map[TopicPartition, OffsetAndMetadata]) = {
+    val (toCommit, retained) =
+      requestedCommitOffsets.partition { case (partition, _) => revoked.contains(partition) }
+    (copy(requestedCommitOffsets = retained), toCommit)
+  }
+
   override def toString: String =
     s"State(partitionGroupState = $partitionGroupState, subscribed = $subscribed, streaming = $streaming)"
 

modules/core/src/test/scala/fs2/kafka/ConsumerSettingsSpec.scala

@@ -289,6 +289,14 @@ final class ConsumerSettingsSpec extends BaseSpec {
         settings.show == settings.toString
       }
     }
+
+    it("should default commitOnRevoke to false and provide withCommitOnRevoke") {
+      assert {
+        !settings.commitOnRevoke &&
+        settings.withCommitOnRevoke(true).commitOnRevoke &&
+        !settings.withCommitOnRevoke(true).withCommitOnRevoke(false).commitOnRevoke
+      }
+    }
   }
 
   val settings =

modules/core/src/test/scala/fs2/kafka/KafkaConsumerSpec.scala

@@ -75,6 +75,51 @@ final class KafkaConsumerSpec extends BaseKafkaSpec {
       }
     }
 
+    it("should commit offsets synchronously on revoke when commitOnRevoke is enabled") {
+      withTopic { topic =>
+        createCustomTopic(topic, partitions = 1)
+        val produced = (0 until 5).map(n => s"key-$n" -> s"value->$n")
+        publishToKafka(topic, produced)
+
+        val settings =
+          consumerSettings[IO]
+            .withGroupId(s"commit-on-revoke-${UUID.randomUUID()}")
+            .withCommitOnRevoke(true)
+
+        // The first consumer processes and commits every record, then shuts down. Closing it
+        // triggers `onPartitionsRevoked`, where the tracked offsets are committed synchronously
+        // (directly on the Java consumer, while the partition is still owned).
+        val consumedFirst =
+          KafkaConsumer
+            .stream(settings)
+            .subscribeTo(topic)
+            .records
+            .evalMap(committable => committable.offset.commit.as(committable.record.key))
+            .take(produced.size.toLong)
+            .compile
+            .toVector
+
+        // A second consumer in the same group must then find nothing left to consume.
+        val consumedSecond =
+          KafkaConsumer
+            .stream(settings)
+            .subscribeTo(topic)
+            .records
+            .map(_.record.key)
+            .interruptAfter(10.seconds)
+            .compile
+            .toVector
+
+        val (first, second) =
+          (for {
+            a <- consumedFirst
+            b <- consumedSecond
+          } yield (a, b)).unsafeRunSync()
+
+        assert(first.size.toLong == produced.size.toLong && second.isEmpty)
+      }
+    }
+
     def testMultipleConsumersCorrectConsumption(
       customizeSettings: ConsumerSettings[IO, String, String] => ConsumerSettings[
         IO,