Add ListenerCallbacks.apply to remove code duplication (#746)

Small refactoring before adding a synchronous listener for rebalance event
akka · Mar 20, 2019 · a8a8268 · a8a8268
1 parent 0f583e8
commit a8a8268
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 36 deletions.
diff --git a/core/src/main/scala/akka/kafka/internal/KafkaConsumerActor.scala b/core/src/main/scala/akka/kafka/internal/KafkaConsumerActor.scala
@@ -26,6 +26,7 @@ import akka.util.JavaDurationConverters._
 import akka.event.LoggingReceive
 import akka.kafka.KafkaConsumerActor.StoppingException
 import akka.kafka._
+import akka.stream.stage.AsyncCallback
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.common.{Metric, MetricName, TopicPartition}
 
@@ -99,6 +100,31 @@ import scala.util.control.NonFatal
   final case class ListenerCallbacks(onAssign: Set[TopicPartition] => Unit, onRevoke: Set[TopicPartition] => Unit)
       extends NoSerializationVerificationNeeded
 
+  object ListenerCallbacks {
+    def apply(subscription: AutoSubscription,
+              sourceActor: ActorRef,
+              partitionAssignedCB: AsyncCallback[Set[TopicPartition]],
+              partitionRevokedCB: AsyncCallback[Set[TopicPartition]]): ListenerCallbacks =
+      KafkaConsumerActor.ListenerCallbacks(
+        assignedTps => {
+          subscription.rebalanceListener.foreach {
+            _.tell(TopicPartitionsAssigned(subscription, assignedTps), sourceActor)
+          }
+          if (assignedTps.nonEmpty) {
+            partitionAssignedCB.invoke(assignedTps)
+          }
+        },
+        revokedTps => {
+          subscription.rebalanceListener.foreach {
+            _.tell(TopicPartitionsRevoked(subscription, revokedTps), sourceActor)
+          }
+          if (revokedTps.nonEmpty) {
+            partitionRevokedCB.invoke(revokedTps)
+          }
+        }
+      )
+  }
+
   /**
    * Copied from the implemented interface:
    *

diff --git a/core/src/main/scala/akka/kafka/internal/SingleSourceLogic.scala b/core/src/main/scala/akka/kafka/internal/SingleSourceLogic.scala
@@ -45,24 +45,7 @@ import scala.concurrent.{Future, Promise}
         log.debug("Revoked partitions: {}. All partitions: {}", revokedTps, tps)
       }
 
-      KafkaConsumerActor.ListenerCallbacks(
-        assignedTps => {
-          autoSubscription.rebalanceListener.foreach {
-            _.tell(TopicPartitionsAssigned(autoSubscription, assignedTps), sourceActor.ref)
-          }
-          if (assignedTps.nonEmpty) {
-            partitionAssignedCB.invoke(assignedTps)
-          }
-        },
-        revokedTps => {
-          autoSubscription.rebalanceListener.foreach {
-            _.tell(TopicPartitionsRevoked(autoSubscription, revokedTps), sourceActor.ref)
-          }
-          if (revokedTps.nonEmpty) {
-            partitionRevokedCB.invoke(revokedTps)
-          }
-        }
-      )
+      KafkaConsumerActor.ListenerCallbacks(autoSubscription, sourceActor.ref, partitionAssignedCB, partitionRevokedCB)
     }
 
     subscription match {

diff --git a/core/src/main/scala/akka/kafka/internal/SubSourceLogic.scala b/core/src/main/scala/akka/kafka/internal/SubSourceLogic.scala
@@ -81,24 +81,7 @@ private abstract class SubSourceLogic[K, V, Msg](
     sourceActor.watch(consumerActor)
 
     def rebalanceListener =
-      KafkaConsumerActor.ListenerCallbacks(
-        assignedTps => {
-          subscription.rebalanceListener.foreach {
-            _.tell(TopicPartitionsAssigned(subscription, assignedTps), sourceActor.ref)
-          }
-          if (assignedTps.nonEmpty) {
-            partitionAssignedCB.invoke(assignedTps)
-          }
-        },
-        revokedTps => {
-          subscription.rebalanceListener.foreach {
-            _.tell(TopicPartitionsRevoked(subscription, revokedTps), sourceActor.ref)
-          }
-          if (revokedTps.nonEmpty) {
-            partitionRevokedCB.invoke(revokedTps)
-          }
-        }
-      )
+      KafkaConsumerActor.ListenerCallbacks(subscription, sourceActor.ref, partitionAssignedCB, partitionRevokedCB)
 
     subscription match {
       case TopicSubscription(topics, _) =>