[greyhound] parallel consumer - interrupt fiber on shutdown timeout (…

…#36642) GitOrigin-RevId: aaccb23bdc18584d42b2240e5b5e63c7d9432838
wix · Sep 8, 2023 · 4de42ee · 4de42ee
1 parent 3c3fd72
commit 4de42ee
Showing 1 changed file with 19 additions and 8 deletions.
diff --git a/core/src/main/scala/com/wixpress/dst/greyhound/core/consumer/EventLoop.scala b/core/src/main/scala/com/wixpress/dst/greyhound/core/consumer/EventLoop.scala
@@ -85,6 +85,7 @@ object EventLoop {
       _                   <- report(CreatingPollOnceFiber(clientId, group, consumerAttributes))
       fiber               <- pollOnce(running, consumer, dispatcher, pausedPartitionsRef, positionsRef, offsets, config, clientId, group, offsetsAndGaps)
                                .repeatWhile(_ == true)
+                               .interruptible
                                .forkDaemon
       _                   <- report(AwaitingPartitionsAssignment(clientId, group, consumerAttributes))
       partitions          <- partitionsAssigned.await
@@ -142,7 +143,9 @@ object EventLoop {
       (joinFiberAndReport(group, clientId, consumerAttributes, fiber).interruptible *>
         shutdownDispatcherAndReport(group, clientId, consumerAttributes, dispatcher))
         .timeout(config.drainTimeout)
-      _       <- ZIO.when(drained.isEmpty)(report(DrainTimeoutExceeded(clientId, group, config.drainTimeout.toMillis, consumerAttributes)))
+      _       <- ZIO.when(drained.isEmpty)(
+                   report(DrainTimeoutExceeded(clientId, group, config.drainTimeout.toMillis, consumerAttributes)) *> fiber.interruptFork
+                 )
       _       <- if (config.consumePartitionInParallel) commitOffsetsAndGaps(consumer, offsetsAndGaps) else commitOffsets(consumer, offsets)
       _       <- report(StoppedEventLoop(clientId, group, consumerAttributes))
     } yield ()
@@ -356,16 +359,18 @@ object EventLoop {
       consumer.commit(committable).catchAll { t => report(FailedToCommitOffsets(t, committable)) *> offsets.update(committable) }
     }
 
-  private def commitOffsetsAndGaps(consumer: Consumer, offsetsAndGaps: OffsetsAndGaps): URIO[GreyhoundMetrics, Unit] =
+  private def commitOffsetsAndGaps(consumer: Consumer, offsetsAndGaps: OffsetsAndGaps): URIO[GreyhoundMetrics, Unit] = {
     offsetsAndGaps.getCommittableAndClear.flatMap { committable =>
       val offsetsAndMetadataToCommit = OffsetsAndGaps.toOffsetsAndMetadata(committable)
-      consumer
-        .commitWithMetadata(offsetsAndMetadataToCommit)
-        .tap(_ => ZIO.when(offsetsAndMetadataToCommit.nonEmpty)(report(CommittedOffsetsAndGaps(committable))))
-        .catchAll { t =>
-          report(FailedToCommitOffsetsAndMetadata(t, offsetsAndMetadataToCommit)) *> offsetsAndGaps.setCommittable(committable)
-        }
+      report(CommittingOffsetsAndGaps(consumer.config.groupId, committable)) *>
+        consumer
+          .commitWithMetadata(offsetsAndMetadataToCommit)
+          .tap(_ => ZIO.when(offsetsAndMetadataToCommit.nonEmpty)(report(CommittedOffsetsAndGaps(committable))))
+          .catchAll { t =>
+            report(FailedToCommitOffsetsAndMetadata(t, offsetsAndMetadataToCommit)) *> offsetsAndGaps.setCommittable(committable)
+          }
     }
+  }
 
   private def commitOffsetsOnRebalance(
     consumer: Consumer,
@@ -460,6 +465,12 @@ object EventLoopMetric {
 
   case class StoppedEventLoop(clientId: ClientId, group: Group, attributes: Map[String, String] = Map.empty) extends EventLoopMetric
 
+  case class CommittingOffsetsAndGaps(
+    groupId: Group,
+    offsetsAndGaps: Map[TopicPartition, OffsetAndGaps],
+    attributes: Map[String, String] = Map.empty
+  ) extends EventLoopMetric
+
   case class JoinedPollOnceFiberBeforeDispatcherShutdown(
     clientId: ClientId,
     group: Group,