Add OnPartitionsLost handler

Author: Arkatufus

src/Akka.Streams.Kafka.Tests/Integration/PlainSourceIntegrationTests.cs

@@ -243,6 +243,13 @@ public void OnRevoke(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions,
                 RevokeEventsCounter.IncrementAndGet();
             }
 
+            /// <inheritdoc />
+            public void OnLost(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions,
+                IRestrictedConsumer consumer)
+            {
+                RevokeEventsCounter.IncrementAndGet();
+            }
+
             /// <inheritdoc />
             public void OnAssign(IImmutableSet<TopicPartition> assignedTopicPartitions, IRestrictedConsumer consumer)
             {

src/Akka.Streams.Kafka/Helpers/PartitionEventHandler.cs

@@ -24,6 +24,11 @@ public interface IPartitionEventHandler
         /// </summary>
         void OnRevoke(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions, IRestrictedConsumer consumer);
 
+        /// <summary>
+        /// Called when partitions are lost
+        /// </summary>
+        void OnLost(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions, IRestrictedConsumer consumer);
+
         /// <summary>
         /// Called when partitions are assigned
         /// </summary>
@@ -50,6 +55,11 @@ public void OnRevoke(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions,
             {
             }
 
+            /// <inheritdoc />
+            public void OnLost(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions, IRestrictedConsumer consumer)
+            {
+            }
+            
             /// <inheritdoc />
             public void OnAssign(IImmutableSet<TopicPartition> assignedTopicPartitions, IRestrictedConsumer consumer)
             {
@@ -82,6 +92,12 @@ public void OnRevoke(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions,
                 _partitionRevokedCallback(revokedTopicPartitions);
             }
 
+            /// <inheritdoc />
+            public void OnLost(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions, IRestrictedConsumer consumer)
+            {
+                OnRevoke(revokedTopicPartitions, consumer);
+            }
+            
             /// <inheritdoc />
             public void OnAssign(IImmutableSet<TopicPartition> assignedTopicPartitions, IRestrictedConsumer consumer)
             {
@@ -115,6 +131,13 @@ public void OnRevoke(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions,
                 _handler2?.OnRevoke(revokedTopicPartitions, consumer);
             }
 
+            /// <inheritdoc />
+            public void OnLost(IImmutableSet<TopicPartitionOffset> revokedTopicPartitions, IRestrictedConsumer consumer)
+            {
+                _handler1?.OnLost(revokedTopicPartitions, consumer);
+                _handler2?.OnLost(revokedTopicPartitions, consumer);
+            }
+            
             /// <inheritdoc />
             public void OnAssign(IImmutableSet<TopicPartition> assignedTopicPartitions, IRestrictedConsumer consumer)
             {

src/Akka.Streams.Kafka/Settings/ConsumerSettings.cs

@@ -394,6 +394,7 @@ private ConsumerSettings<TKey, TValue> Copy(
         public Confluent.Kafka.IConsumer<TKey, TValue> CreateKafkaConsumer(Action<IConsumer<TKey, TValue>, Error> consumeErrorHandler = null,
                                                                            Action<IConsumer<TKey, TValue>, List<TopicPartition>> partitionAssignedHandler = null,
                                                                            Action<IConsumer<TKey, TValue>, List<TopicPartitionOffset>> partitionRevokedHandler = null,
+                                                                           Action<IConsumer<TKey, TValue>, List<TopicPartitionOffset>> partitionLostHandler = null,
                                                                            Action<IConsumer<TKey, TValue>, string> statisticHandler = null)
         {
             RebalanceListener = new RebalanceListener<TKey, TValue>(
@@ -409,6 +410,7 @@ public Confluent.Kafka.IConsumer<TKey, TValue> CreateKafkaConsumer(Action<IConsu
                 .SetErrorHandler((c, e) => consumeErrorHandler?.Invoke(c, e))
                 .SetPartitionsAssignedHandler((c, partitions) => partitionAssignedHandler?.Invoke(c, partitions))
                 .SetPartitionsRevokedHandler((c, partitions) => partitionRevokedHandler?.Invoke(c, partitions))
+                .SetPartitionsLostHandler((c, partitions) => partitionLostHandler?.Invoke(c, partitions))
                 .SetStatisticsHandler((c, json) => statisticHandler?.Invoke(c, json))
                 .Build();
         }

src/Akka.Streams.Kafka/Stages/Consumers/Actors/KafkaConsumerActor.cs

@@ -147,6 +147,17 @@ private void PartitionsRevokedHandler(IImmutableSet<TopicPartitionOffset> partit
             _rebalanceInProgress.GetAndSet(true);
         }
 
+        // This is RebalanceListener.OnPartitionsLost on JVM
+        private void PartitionsLostHandler(IImmutableSet<TopicPartitionOffset> partitions)
+        {
+            var watch = Stopwatch.StartNew();
+            _partitionEventHandler.OnLost(partitions, _restrictedConsumer);
+            watch.Stop();
+            CheckDuration(watch, "onLost");
+            
+            _commitRefreshing.Revoke(partitions.Select(tp => tp.TopicPartition).ToImmutableHashSet());
+        }
+
         private void RebalancePostStop()
         {
             var currentTopicPartitions = _consumer.Assignment.ToImmutableList();
@@ -490,17 +501,15 @@ private void Poll()
 
                 using (var cts = new CancellationTokenSource(_settings.PollTimeout))
                 {
-                    var (polled, exception) = PollKafka(cts.Token);
                     try
                     {
+                        var polled = PollKafka(cts.Token);
                         ProcessResult(partitionsToFetch, polled);
                     }
                     catch (Exception e)
                     {
                         ProcessExceptions(e);
                     }
-
-                    ProcessExceptions(exception);
                 }
             }
 
@@ -529,29 +538,20 @@ private void ProcessExceptions(Exception exception)
             Context.Stop(Self);
         }
 
-        private (List<ConsumeResult<K, V>>, Exception) PollKafka(CancellationToken token)
+        private List<ConsumeResult<K, V>> PollKafka(CancellationToken token)
         {
             ConsumeResult<K, V> consumed = null;
             var i = 10; // 10 poll attempts
             var timeout = Math.Max((int) _pollTimeout.TotalMilliseconds / i, 1);
             var polled = new List<ConsumeResult<K, V>>();
             do
             {
-                try
-                {
-                    // this would return immediately if there are messages waiting inside the client queue buffer
-                    consumed = _consumer.Consume(timeout);
-                }
-                catch (Exception e)
-                {
-                    return (polled, e);
-                }
-                if (consumed != null)
-                    polled.Add(consumed);
+                // this would return immediately if there are messages waiting inside the client queue buffer
+                consumed = _consumer.Consume(timeout);
                 i--;
             } while (i > 0 && consumed != null && !token.IsCancellationRequested);
 
-            return (polled, null);
+            return polled;
         }
 
         private void ProcessResult(IImmutableSet<TopicPartition> partitionsToFetch, List<ConsumeResult<K,V>> rawResult)