rebalance finished event

src/KafkaNET.Library/Consumers/ZookeeperConsumerConnector.cs

@@ -21,6 +21,7 @@ namespace Kafka.Client.Consumers
     using Kafka.Client.Serialization;
     using Kafka.Client.Utils;
     using Kafka.Client.ZooKeeperIntegration;
+    using Kafka.Client.ZooKeeperIntegration.Events;
     using Kafka.Client.ZooKeeperIntegration.Listeners;
     using System;
     using System.Collections.Concurrent;
@@ -33,6 +34,7 @@ namespace Kafka.Client.Consumers
     /// </summary>
     public class ZookeeperConsumerConnector : KafkaClientBase, IZookeeperConsumerConnector
     {
+        private readonly EventHandler<ConsumerRebalanceEventArgs> consumerRebalanceFinishedHandler;
         public static log4net.ILog Logger = log4net.LogManager.GetLogger(typeof(ZookeeperConsumerConnector));
         public static readonly int MaxNRetries = 4;
         public static bool UseSharedStaticZookeeperClient = true;
@@ -74,7 +76,8 @@ public ZookeeperConsumerConnector(ConsumerConfiguration config,
                                           bool enableFetcher,
                                           EventHandler rebalanceHandler = null,
                                           EventHandler zkDisconnectedHandler = null,
-                                          EventHandler zkExpiredHandler = null)
+                                          EventHandler zkExpiredHandler = null,
+                                          EventHandler<ConsumerRebalanceEventArgs> rebalanceFinishedHandler = null)
         {
             if (string.IsNullOrEmpty(config.GroupId))
             {
@@ -94,6 +97,7 @@ public ZookeeperConsumerConnector(ConsumerConfiguration config,
                 this.consumerRebalanceHandler = rebalanceHandler;
                 this.zkSessionDisconnectedHandler = zkDisconnectedHandler;
                 this.zkSessionExpiredHandler = zkExpiredHandler;
+                this.consumerRebalanceFinishedHandler = rebalanceFinishedHandler;
 
                 if (this.config.AutoCommit)
                 {
@@ -596,6 +600,10 @@ private IDictionary<string, IList<KafkaMessageStream<TData>>> Consume<TData>(IDi
             {
                 loadBalancerListener.ConsumerRebalance += this.consumerRebalanceHandler;
             }
+            if (this.consumerRebalanceFinishedHandler != null)
+            {
+                loadBalancerListener.ConsumerRebalanceFinished += this.consumerRebalanceFinishedHandler;
+            }
 
             stopAsyncRebalancing.Add(loadBalancerListener.StopRebalance);
             this.RegisterConsumerInZk(dirs, consumerIdString, topicCount);

src/KafkaNET.Library/KafkaNET.Library.csproj

@@ -174,6 +174,7 @@
     <Compile Include="Consumers\IKafkaMessageStream.cs" />
     <Compile Include="Consumers\IZookeeperConsumerConnector.cs" />
     <Compile Include="ZooKeeperIntegration\Events\ChildChangedEventItem.cs" />
+    <Compile Include="ZooKeeperIntegration\Events\ConsumerRebalanceEventArgs.cs" />
     <Compile Include="ZooKeeperIntegration\Events\DataChangedEventItem.cs" />
     <Compile Include="ZooKeeperIntegration\Events\ZooKeeperChildChangedEventArgs.cs" />
     <Compile Include="ZooKeeperIntegration\Events\ZooKeeperDataChangedEventArgs.cs" />

src/KafkaNET.Library/ZooKeeperIntegration/Events/ConsumerRebalanceEventArgs.cs

@@ -0,0 +1,15 @@
+using System;
+using System.Collections.Generic;
+
+namespace Kafka.Client.ZooKeeperIntegration.Events
+{
+    public class ConsumerRebalanceEventArgs : EventArgs
+    {
+        public IDictionary<string, ICollection<int>> TopicPartitionMap { get; private set; }
+
+        public ConsumerRebalanceEventArgs(Dictionary<string, ICollection<int>> topicPartitionMap)
+        {
+            TopicPartitionMap = topicPartitionMap;
+        }
+    }
+}

src/KafkaNET.Library/ZooKeeperIntegration/Listeners/ZKRebalancerListener.cs

@@ -42,6 +42,7 @@ internal class ZKRebalancerListener<TData> : IZooKeeperChildListener
         public static log4net.ILog Logger = log4net.LogManager.GetLogger("ZKRebalancerListener");
 
         public event EventHandler ConsumerRebalance;
+        public event EventHandler<ConsumerRebalanceEventArgs> ConsumerRebalanceFinished;
 
         private IDictionary<string, IList<string>> oldPartitionsPerTopicMap = new Dictionary<string, IList<string>>();
         private IDictionary<string, IList<string>> oldConsumersPerTopicMap = new Dictionary<string, IList<string>>();
@@ -206,6 +207,18 @@ protected virtual void OnConsumerRebalance(EventArgs args)
             }
         }
 
+        protected virtual void OnConsumerRebalanceFinished(ConsumerRebalanceEventArgs args)
+        {
+            try
+            {
+                ConsumerRebalanceFinished?.Invoke(this, args);
+            }
+            catch (Exception ex)
+            {
+                Logger.Error("Exception occurred within event handler for ConsumerRebalanceFinished event: " + ex.Message);
+            }
+        }
+
         private void SyncedRebalance(CancellationTokenSource cancellationTokenSource)
         {
             Logger.InfoFormat("Consumer {0} has entered rebalance", consumerIdString);
@@ -252,6 +265,8 @@ private void SyncedRebalance(CancellationTokenSource cancellationTokenSource)
                 isRebalanceRunning = false;
             }
 
+            OnConsumerRebalanceFinished(new ConsumerRebalanceEventArgs(topicRegistry.ToDictionary(t => t.Key, t => t.Value.Keys)));
+
             Logger.InfoFormat("Consumer {0} has exited rebalance", consumerIdString);
         }